diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index 023b51a156..ecee81635e 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -36,7 +36,7 @@ jobs: platform_id: macosx_x86_64 dp_variant: cpu # macos-arm64 - - os: macos-14 + - os: macos-15 python: 311 platform_id: macosx_arm64 dp_variant: cpu @@ -61,7 +61,7 @@ jobs: version: "0.12.0" if: runner.os != 'Linux' - name: Build wheels - uses: pypa/cibuildwheel@v4.2 + uses: pypa/cibuildwheel@v4.3 env: CIBW_BUILD_VERBOSITY: 1 CIBW_ARCHS: all diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml index 13e3e51a32..361b709704 100644 --- a/.github/workflows/test_cc.yml +++ b/.github/workflows/test_cc.yml @@ -94,6 +94,7 @@ jobs: TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp + DPA4SPIN_LMP_REQUIRED: ${{ matrix.enable_pytorch && '1' || '0' }} ENABLE_TENSORFLOW: ${{ matrix.enable_tensorflow && '1' || '0' }} ENABLE_PYTORCH: ${{ matrix.enable_pytorch && '1' || '0' }} ENABLE_JAX: ${{ matrix.enable_tensorflow && '1' || '0' }} diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 01f8dcfa76..c72d01f06a 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -20,9 +20,19 @@ env: DP_ENABLE_PADDLE: "0" jobs: test_python: - name: Test Python on CUDA + name: Test Python on CUDA (${{ matrix.shard }}) runs-on: gpu - # The serial Python test suite can take several hours on the GPU runner. + strategy: + fail-fast: false + matrix: + include: + - shard: pt_expt-model + pytest-args: source/tests/pt_expt/model + - shard: pt_expt-other + pytest-args: source/tests/pt_expt --ignore=source/tests/pt_expt/model + - shard: other + pytest-args: source/tests --ignore=source/tests/pt_expt + # Each shard has its own CUDA build and can take several hours. timeout-minutes: 480 # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845 # container: @@ -71,7 +81,7 @@ jobs: - name: Verify float64 CUDA kernels are enabled run: | python -c "import torch; import deepmd.pt.cxx_op; from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import op_available, float64_available; assert torch.cuda.is_available(); assert op_available(); assert float64_available() is True" - - run: python -m pytest source/tests --ignore=source/tests/pd + - run: python -m pytest ${{ matrix.pytest-args }} --ignore=source/tests/pd env: NUM_WORKERS: 0 CUDA_VISIBLE_DEVICES: 0 @@ -81,8 +91,16 @@ jobs: FLAGS_use_stride_compute_kernel: 0 test_cc: - name: Test C++ on CUDA + name: Test ${{ matrix.name }} on CUDA runs-on: gpu + strategy: + fail-fast: false + matrix: + include: + - suite: main + name: C++ + - suite: kokkos + name: Kokkos if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group') steps: # Jobs run on separate runners, so the C++ job needs its own complete @@ -92,7 +110,7 @@ jobs: with: python-version: "3.11" # cache: 'pip' - - name: Install wget and unzip + - name: Install build prerequisites run: sudo apt-get update && sudo apt-get install -y wget unzip - uses: lukka/get-cmake@latest with: @@ -115,7 +133,15 @@ jobs: - run: | export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])') 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 -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit + # Use the LAMMPS wheel's MPI distribution for every build and launch. + # Ubuntu 24.04 MPICH can start independent singleton worlds: + # https://github.com/mpi4py/mpi4py/issues/547 + source/install/uv_with_retry.sh pip install --system mpich + export MPI_HOME=$(python -c 'import sys; print(sys.prefix)') + export MPICC=$MPI_HOME/bin/mpicc + export LD_LIBRARY_PATH=$MPI_HOME/lib:${LD_LIBRARY_PATH:-} + echo "MPI_HOME=$MPI_HOME" >> "$GITHUB_ENV" + source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit --reinstall-package mpi4py --no-binary mpi4py # See https://github.com/jax-ml/jax/issues/29042 source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13' env: @@ -123,6 +149,26 @@ jobs: DP_ENABLE_NATIVE_OPTIMIZATION: 1 DP_ENABLE_PYTORCH: 1 - run: dp --version + - name: Verify the two-rank MPI runtime + run: | + export LD_LIBRARY_PATH=$MPI_HOME/lib:${LD_LIBRARY_PATH:-} + timeout 60 "$MPI_HOME/bin/mpiexec" -n 2 python -c ' + import os + from pathlib import Path + from mpi4py import MPI + comm = MPI.COMM_WORLD + assert comm.size == 2, MPI.Get_library_version() + assert comm.allgather(comm.rank) == [0, 1] + libraries = { + Path(line.split()[-1]).resolve() + for line in Path("/proc/self/maps").read_text().splitlines() + if Path(line.split()[-1]).name.startswith("libmpi") + } + expected = (Path(os.environ["MPI_HOME"]) / "lib").resolve() + assert libraries and all(p.parent == expected for p in libraries), libraries + if comm.rank == 0: + print(MPI.Get_library_version(), libraries) + ' - name: Test the default CUDA dtype contract run: | python -c "import torch; import deepmd.pt.cxx_op; from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import op_available, float64_available; assert torch.cuda.is_available(); assert op_available(); assert float64_available() is False" @@ -131,9 +177,12 @@ jobs: CUDA_VISIBLE_DEVICES: 0 OMP_NUM_THREADS: 1 - name: Convert models + if: matrix.suite == 'main' run: source/tests/infer/convert-models.sh - - run: | - export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH + - name: Build and test C++ + if: matrix.suite == 'main' + run: | + export LD_LIBRARY_PATH=$MPI_HOME/lib:$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH source/install/test_cc_local.sh env: OMP_NUM_THREADS: 1 @@ -141,10 +190,37 @@ jobs: TF_INTER_OP_PARALLELISM_THREADS: 1 CMAKE_GENERATOR: Ninja DP_VARIANT: cuda - DP_USE_MPICH2: 1 ENABLE_PADDLE: FALSE - - run: | - export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$GITHUB_WORKSPACE/dp_test/lib:$LD_LIBRARY_PATH + - name: Prepare the Kokkos C API and long-range fixtures + if: matrix.suite == 'kokkos' + run: | + export LD_LIBRARY_PATH=$MPI_HOME/lib:$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:${LD_LIBRARY_PATH:-} + source/install/test_cc_local.sh --build-only + source/install/test_cc_local.sh --models-only=kokkos + env: + CMAKE_GENERATOR: Ninja + DP_VARIANT: cuda + ENABLE_PADDLE: FALSE + CUDA_VISIBLE_DEVICES: 0 + OMP_NUM_THREADS: 1 + TF_INTRA_OP_PARALLELISM_THREADS: 1 + TF_INTER_OP_PARALLELISM_THREADS: 1 + - name: Test CUDA Kokkos LAMMPS + if: matrix.suite == 'kokkos' + run: | + export LD_LIBRARY_PATH=$MPI_HOME/lib:$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:${LD_LIBRARY_PATH:-} + source/install/test_lammps_kokkos.sh + env: + CMAKE_GENERATOR: Ninja + DP_VARIANT: cuda + CUDA_VISIBLE_DEVICES: 0 + OMP_NUM_THREADS: 1 + TF_INTRA_OP_PARALLELISM_THREADS: 1 + TF_INTER_OP_PARALLELISM_THREADS: 1 + - name: Test LAMMPS and i-PI + if: matrix.suite == 'main' + run: | + export LD_LIBRARY_PATH=$MPI_HOME/lib:$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$GITHUB_WORKSPACE/dp_test/lib:$LD_LIBRARY_PATH export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH python -m pytest -s source/lmp/tests || (cat log.lammps && exit 1) python -m pytest source/ipi/tests @@ -153,6 +229,7 @@ jobs: TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp + DPA4SPIN_LMP_REQUIRED: 1 CUDA_VISIBLE_DEVICES: 0 ENABLE_PADDLE: 0 pass: @@ -165,4 +242,5 @@ jobs: uses: re-actors/alls-green@release/v1 with: jobs: ${{ toJSON(needs) }} - allowed-skips: test_python, test_cc + # Require all Python, C++, and Kokkos matrix jobs on requested GPU runs. + allowed-skips: ${{ !(github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')) && 'test_python, test_cc' || '' }} diff --git a/backend/find_pytorch.py b/backend/find_pytorch.py index f969c5ebce..1e7117b54a 100644 --- a/backend/find_pytorch.py +++ b/backend/find_pytorch.py @@ -126,6 +126,19 @@ def get_pt_requirement(pt_version: str = "") -> dict: else: mpi_requirement = [] + # Vesin's torch bindings require PyTorch 2.3. Older PyTorch targets and + # macOS x86_64, whose last PyTorch wheels are 2.2, use native neighbors. + vesin_requirement = [] + if pt_version == "" or Version(pt_version) >= Version("2.3"): + supported_platform = ( + "platform_system != 'Darwin' or platform_machine != 'x86_64'" + ) + vesin_requirement = [ + f"vesin[torch]>=0.6; {supported_platform}", + # The extra does not constrain the independently versioned bindings. + f"vesin-torch>=0.6; {supported_platform}", + ] + return { "torch": [ # uv has different local version behaviors, i.e. `==2.3.1` cannot match `==2.3.1+cpu` @@ -137,11 +150,7 @@ def get_pt_requirement(pt_version: str = "") -> dict: # https://github.com/pytorch/pytorch/commit/7e0c26d4d80d6602aed95cb680dfc09c9ce533bc else "torch>=2.1.0", "e3nn>=0.5.9", - # O(N) cell-list neighbor list for fast Python/ASE inference; the - # torch bindings (vesin-torch) ship only as a PyPI extra, so keep it - # under the torch extra rather than the core deps (conda-forge has - # vesin but not vesin-torch). - "vesin[torch]", + *vesin_requirement, # GPU O(N) cell-list neighbor list for large systems. Restricted to # Python >= 3.11 (the package requires it while deepmd-kit still # supports 3.10) and to Linux: it is a CUDA package, and its diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 0eca4776ed..402e09ba77 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -1,9 +1,24 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """ASE calculator interface module.""" +import logging +import math +import weakref +from contextlib import ( + ExitStack, +) +from functools import ( + partial, +) from pathlib import ( Path, ) +from shutil import ( + copyfile, +) +from tempfile import ( + TemporaryDirectory, +) from typing import ( TYPE_CHECKING, Any, @@ -12,6 +27,7 @@ Union, ) +import numpy as np from ase.calculators.calculator import ( Calculator, PropertyNotImplementedError, @@ -23,6 +39,10 @@ ) if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + from ase import ( Atoms, ) @@ -30,8 +50,84 @@ NeighborList, ) + from deepmd.infer.deep_eval import ( + DeepEvalBackend, + ) + __all__ = ["DP"] +log = logging.getLogger(__name__) + + +def _compress_dpa4c_for_ase(model: Any) -> bool: + """Select supported DPA4C compression without changing other model families.""" + from deepmd.pt_expt.descriptor.dpa4c import ( + DescrptDPA4C, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + compression_ineligibility_reason, + ) + + parts = model.atomic_model.fused_decomposition() + if ( + parts is None + or parts[0] is None + or not isinstance(parts[0].descriptor, DescrptDPA4C) + ): + return False + reason = compression_ineligibility_reason(parts[0].descriptor) + if reason is None: + return True + log.warning( + "DPA4C compression is unavailable: %s Freezing without compression.", reason + ) + return False + + +def _checkpoint_freezer( + evaluator: "DeepEvalBackend", head: str | None, *, freeze_requested: bool +) -> "Callable[..., None] | None": + """Select an exporter for the loaded checkpoint's model and head.""" + from deepmd.pt.infer.deep_eval import DeepEval as PTDeepEval + from deepmd.pt.model.model.sezm_model import ( + SeZMModel, + ) + from deepmd.pt_expt.infer.deep_eval import DeepEval as PTExptDeepEval + + if isinstance(evaluator, PTDeepEval) and isinstance( + evaluator.get_model(), SeZMModel + ): + from deepmd.pt.entrypoints.freeze_pt2 import ( + freeze_sezm_to_pt2, + ) + + return partial( + freeze_sezm_to_pt2, + head=head if evaluator.multi_task else None, + with_comm=False, + ) + if ( + isinstance(evaluator, PTExptDeepEval) + and evaluator.metadata["lower_input_kind"] == "graph" + ): + from deepmd.pt_expt.entrypoints.main import ( + freeze, + ) + + # The checkpoint loader selects Default when no head is supplied. + return partial( + freeze, + head="Default" if head is None else head, + lower_kind="graph", + with_comm=False, + compress=( + _compress_dpa4c_for_ase(evaluator._dpmodel) + if freeze_requested + else False + ), + ) + return None + class DP(Calculator): """Implementation of ASE deepmd calculator. @@ -51,6 +147,34 @@ class DP(Calculator): The neighbor list object. If None, then build the native neighbor list. head : Union[str, None], optional a specific model branch choosing from pretrained model, by default None + nlist_backend : str, optional + Neighbor search for models using the extended-coordinate input schema. + The default ``"auto"`` selects an available accelerated search. + neighbor_graph_method : str, optional + Graph search for PT-expt graph models or PT SeZM/DPA4, by default + ``"auto"``. Single-frame CUDA inference prefers Vesin. PT-expt graph + models also accept ``"dense"``, ``"ase"``, ``"cell"``, ``"vesin"``, + and ``"nv"``; SeZM edge-vector inputs accept ``"vesin"`` and ``"nv"``. + neighbor_skin : float, optional + Vesin candidate-search skin in angstrom, by default 1.0. Set to zero + to rebuild candidates every step. All searches still use the model's + exact cutoff for the returned neighbors. Applies to both graph and + extended-coordinate inference, independently of checkpoint freezing. + freeze : bool, optional + Export an eligible ``.pt`` checkpoint to ``.pt2`` for inference, + by default False. + Supports PT DPA4/SeZM and graph-eligible PT-expt models. Each call + retains a private artifact for its evaluator's lifetime and atomically + publishes a copy as ``-ase_single.pt2`` beside the + checkpoint. Concurrent processes retain their own selected model heads; + the last publication replaces the shared copy. Export uses the + backend's freeze-time kernel policy and omits the + multi-rank communication artifact. DPA4C models are also compressed + when their configuration supports it; otherwise a warning explains + why the export remains uncompressed. Other models ignore this option + and are loaded directly. Direct checkpoint inference for supported + models can be substantially slower; False is intended for debugging + without the compilation overhead. Examples -------- @@ -92,21 +216,74 @@ def __init__( neighbor_list: Optional["NeighborList"] = None, head: str | None = None, nlist_backend: str = "auto", + freeze: bool = False, + neighbor_graph_method: str = "auto", + neighbor_skin: float = 1.0, **kwargs: Any, ) -> None: Calculator.__init__(self, label=label, **kwargs) + if not math.isfinite(neighbor_skin) or neighbor_skin < 0: + raise ValueError( + f"neighbor_skin must be finite and non-negative, got {neighbor_skin}" + ) + model_path = Path(model).resolve() self.dp = DeepPot( - str(Path(model).resolve()), + str(model_path), neighbor_list=neighbor_list, head=head, nlist_backend=nlist_backend, + neighbor_graph_method=neighbor_graph_method, + neighbor_skin=neighbor_skin, ) + freezer = ( + _checkpoint_freezer(self.dp.deep_eval, head, freeze_requested=freeze) + if model_path.suffix == ".pt" + else None + ) + if freezer is not None and freeze: + output = model_path.with_name(f"{model_path.stem}-ase_single.pt2") + # Release the eager model before the exporter constructs its model. + self.dp.close() + del self.dp + log.info("Freezing ASE checkpoint %s to %s", model_path, output) + with ExitStack() as cleanup: + temporary = TemporaryDirectory( + prefix=f".{model_path.stem}-ase-", dir=model_path.parent + ) + cleanup.callback(temporary.cleanup) + private = Path(temporary.name) / output.name + freezer(str(model_path), str(private)) + self.dp = DeepPot( + str(private), + neighbor_list=neighbor_list, + nlist_backend=nlist_backend, + neighbor_graph_method=neighbor_graph_method, + neighbor_skin=neighbor_skin, + ) + # A separate inode keeps later writes to the public artifact + # from changing the file reopened by evaluator serialization. + publication = Path(temporary.name) / "publish.pt2" + copyfile(private, publication) + publication.replace(output) + weakref.finalize(self.dp.deep_eval, temporary.cleanup) + cleanup.pop_all() + elif freezer is not None: + log.warning( + "ASE inference directly from a .pt checkpoint can be SUBSTANTIALLY " + "slower than a frozen .pt2 model. Use DP(..., freeze=True), or " + "freeze the checkpoint manually and load the .pt2 file. " + "freeze=False is intended for debugging and similar short runs " + "that avoid the upfront compilation overhead." + ) if type_dict: self.type_dict = type_dict else: self.type_dict = dict( zip(self.dp.get_type_map(), range(self.dp.get_ntypes()), strict=True) ) + self._cached_numbers: np.ndarray | None = None + self._cached_types: np.ndarray | None = None + self._cached_type_dict: dict[str, int] = {} def calculate( self, @@ -134,8 +311,22 @@ def calculate( cell = self.atoms.get_cell().reshape([1, -1]) else: cell = None - symbols = self.atoms.get_chemical_symbols() - atype = [self.type_dict[k] for k in symbols] + numbers = self.atoms.numbers + if ( + not np.array_equal(numbers, self._cached_numbers) + or self.type_dict != self._cached_type_dict + ): + self._cached_types = np.fromiter( + ( + self.type_dict[symbol] + for symbol in self.atoms.get_chemical_symbols() + ), + dtype=np.int32, + count=len(numbers), + ) + self._cached_numbers = numbers.copy() + self._cached_type_dict = self.type_dict.copy() + atype = self._cached_types fparam = self.atoms.info.get("fparam", None) aparam = self.atoms.info.get("aparam", None) diff --git a/deepmd/common.py b/deepmd/common.py index fea637328a..94d07e3ab0 100644 --- a/deepmd/common.py +++ b/deepmd/common.py @@ -29,8 +29,9 @@ from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) -from deepmd.utils.path import ( - DPPath, +from deepmd.utils.data_path import ( + expand_sys_str, + rglob_sys_str, ) __all__ = [ @@ -41,6 +42,7 @@ "get_np_precision", "j_loader", "make_default_mesh", + "rglob_sys_str", "select_idx_map", ] @@ -188,50 +190,6 @@ def j_loader(filename: str | Path) -> dict[str, Any]: raise TypeError("config file must be json, or yaml/yml") -def expand_sys_str(root_dir: str | Path) -> list[str]: - """Recursively iterate over directories taking those that contain `type.raw` file. - - Parameters - ---------- - root_dir : Union[str, Path] - starting directory - - Returns - ------- - list[str] - list of string pointing to system directories - """ - root_dir = DPPath(root_dir) - matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()] - if (root_dir / "type.raw").is_file(): - matches.append(str(root_dir)) - return matches - - -def rglob_sys_str(root_dir: str, patterns: list[str]) -> list[str]: - """Recursively iterate over directories taking those that contain `type.raw` file. - - Parameters - ---------- - root_dir : str, Path - starting directory - patterns : list[str] - list of glob patterns to match directories - - Returns - ------- - list[str] - list of string pointing to system directories - """ - root_dir = Path(root_dir) - matches = [] - for pattern in patterns: - matches.extend( - [str(d) for d in root_dir.rglob(pattern) if (d / "type.raw").is_file()] - ) - return list(set(matches)) # remove duplicates - - def get_np_precision(precision: "_PRECISION") -> np.dtype: """Get numpy precision constant from string. diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 25b770a5fb..a330c6de79 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -51,6 +51,7 @@ check_preset_out_bias, normalize_preset_out_bias, preset_out_bias_rows, + recover_nonfinite_preset_bias, remap_preset_out_bias, ) @@ -221,6 +222,67 @@ def fused_decomposition(self) -> "tuple[Any, Any] | None": """ return None + def has_long_range(self) -> bool: + """Whether the atomic energy carries a long-range module beside a fitting. + + Generic capability (concrete default ``False``): a model wrapping the + atomic model asks it whether the energy holds a long-range interaction, + e.g. before combining it with a scheme defined for a learned local + energy alone. + """ + return False + + def required_topologies(self) -> tuple[str, ...]: + """Names of the topologies the graph forward reads beside the neighbor graph. + + Generic capability (concrete default: none) in the sense of + :meth:`uses_graph_lower`: a long-range module, whose interaction has no + cutoff, declares the topology it is evaluated on (see + :func:`~deepmd.dpmodel.utils.neighbor_graph.build_topologies`), and + every caller of the graph forward builds what is declared. + """ + return () + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of :meth:`required_topologies`. + + Generic capability (concrete default: none) in the sense of + :meth:`required_topologies`: a module whose topology depends on its + own state, such as the wave-number cutoff of a periodic long-range + module, which follows its kernel, gives it here, by topology name. + """ + return {} + + def build_topologies( + self, + coord: Any, + n_node: Any, + box: Any, + *, + atype: Any, + node_mask: Any = None, + pair_excl: Any = None, + ) -> dict[str, Any]: + """Build the topologies of :meth:`required_topologies` for one batch. + + See :func:`~deepmd.dpmodel.utils.neighbor_graph.build_topologies`; + the build parameters are those of :meth:`topology_params`. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_topologies, + ) + + return build_topologies( + self.required_topologies(), + coord, + n_node, + box, + node_mask, + atype=atype, + pair_excl=pair_excl, + topology_params=self.topology_params(), + ) + def supports_native_spin(self) -> bool: """Returns whether this atomic model consumes a per-atom spin input. @@ -493,6 +555,8 @@ def forward_common_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict: """Graph analogue of :meth:`forward_common_atomic` on the flat node axis. @@ -529,6 +593,13 @@ def forward_common_atomic_graph( (and, from there, the descriptor's ``call_graph``). ``None`` for non-parallel inference (default). Mirrors :meth:`forward_common_atomic`'s ``comm_dict`` on the dense route. + topologies + the topologies of :meth:`required_topologies` by name, built on the + node axis of ``graph``; ``None`` when none is required. + node_mask + flat mask of the nodes an enclosing model counts (real, owned and + not excluded there), combined with this model's own; ``None`` + imposes nothing further. Returns ------- @@ -537,6 +608,8 @@ def forward_common_atomic_graph( """ graph, atype_clamped, output_mask = self._prepare_graph_inputs(graph, atype) + if node_mask is not None: + output_mask = output_mask & node_mask ret_dict = self.forward_atomic_graph( graph, atype_clamped, @@ -545,6 +618,8 @@ def forward_common_atomic_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, + node_mask=output_mask, ) return self._finalize_atomic_ret(ret_dict, output_mask, atype) @@ -822,6 +897,7 @@ def change_out_bias( stat_file_path: DPPath | None = None, bias_adjust_mode: str = "change-by-statistic", observed_type: list[str] | None = None, + preset_out_bias: dict | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -845,6 +921,14 @@ def change_out_bias( ('set-by-statistic') or their stored bias ('change-by-statistic'); no statistics are computed, read or written for such an output, whose output std keeps its stored value. + preset_out_bias : dict, optional + An explicit downstream preset. Assigned outputs use this bias and + skip fitting; other outputs fit residuals against the stored bias. + The model preset is used when omitted, preserving training and + fine-tuning behavior. Pass an empty dict to fit from the data + without reapplying the model's stored preset. + Before residual fitting, non-finite stored cells covered by the + checkpoint preset are recovered; finite adjusted cells are retained. stat_file_path : Optional[DPPath] The path to the stat file. observed_type : list[str], optional @@ -859,17 +943,32 @@ def change_out_bias( if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"): raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) change = bias_adjust_mode == "change-by-statistic" + preset = ( + normalize_preset_out_bias(preset_out_bias, self.type_map) + if preset_out_bias is not None + else self.preset_out_bias + ) distinguish_types = self.get_compute_stats_distinguish_types() - check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types) + check_preset_out_bias(preset, self.bias_keys, distinguish_types) sampled = sample_merged + # A per-call override must not propagate a corrupt checkpoint baseline. + if change and preset_out_bias is not None and self.preset_out_bias: + recovered = recover_nonfinite_preset_bias( + self.preset_out_bias, + self.type_map, + to_numpy_array(self.out_bias), + self.bias_keys, + [self.atomic_output_def()[kk].size for kk in self.bias_keys], + ) + self._store_out_stat(recovered) # === Step 1. Outputs fixed by the preset === fitted_keys = self.bias_keys - if self.preset_out_bias: + if preset: if observed_type is None: sampled = sample_merged() if callable(sample_merged) else sample_merged observed_type = collect_observed_types(sampled, self.type_map) rows = preset_out_bias_rows( - self.preset_out_bias, + preset, self.type_map, observed_type, to_numpy_array(self.out_bias), @@ -989,9 +1088,18 @@ def model_forward( self.get_rcut(), pair_excl=self.pair_excl, ) + atype_flat = xp.reshape(atype, (-1,)) + topologies = self.build_topologies( + xp.reshape(coord, (-1, 3)), + graph.n_node, + box, + atype=atype_flat, + node_mask=atype_flat >= 0, + pair_excl=self.pair_excl, + ) atomic_ret = self.forward_common_atomic_graph( graph, - xp.reshape(atype, (-1,)), + atype_flat, fparam=fparam, aparam=( xp.reshape( @@ -1003,6 +1111,7 @@ def model_forward( ), charge_spin=charge_spin, spin=None if spin is None else xp.reshape(spin, (-1, 3)), + topologies=topologies, ) # The graph route works on a flat node axis; restore the # per-frame layout the dense route returns. diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index b9f20ed8c7..b92deb413b 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -24,6 +24,10 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) +from deepmd.dpmodel.long_range import ( + BaseLongRange, + check_long_range_support, +) from deepmd.dpmodel.output_def import ( FittingOutputDef, ) @@ -102,6 +106,12 @@ class DPAtomicModel(BaseAtomicModel): :math:`\mathcal{F}` is the fitting network, and :math:`\mathbf{y}^i` is the predicted atomic property (energy, dipole, etc.). + An energy model may carry a long-range module beside its fitting: it reads + the same descriptor and adds its own atomic energy to the fitting output, + so the atomic energy :math:`E^i` is the complete potential, including the + interactions without a cutoff. The module is evaluated on the + NeighborGraph route. + Parameters ---------- descriptor @@ -111,6 +121,8 @@ class DPAtomicModel(BaseAtomicModel): type_map Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. + long_range + Long-range module added to the energy output of the fitting. """ @@ -121,12 +133,14 @@ def __init__( descriptor: BaseDescriptor, fitting: BaseFitting, type_map: list[str], + long_range: BaseLongRange | None = None, **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) self.descriptor = descriptor self.fitting_net = fitting self.fitting_net.reinit_exclude(self.atom_exclude_types) + self.long_range = long_range self.type_map = type_map self.add_chg_spin_ebd: bool = self.descriptor.has_chg_spin_ebd() # Structural capability: only descriptors with a native spin @@ -148,6 +162,8 @@ def __init__( self.add_spin_ebd: bool = ( self._supports_native_spin and self.descriptor.use_spin is not None ) + if self.long_range is not None: + self._check_long_range() super().init_out_stat() # === Reference of the isolated atoms === # The fitting references its output only when the preset fixes the @@ -155,6 +171,43 @@ def __init__( resolve_vacuum_ref(self.fitting_net, self.preset_out_bias) self.init_vacuum_conditions() + def _check_long_range(self) -> None: + """Validate that the long-range module can join this descriptor and fitting.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + ) + + check_long_range_support(self.descriptor) + if not self.descriptor.uses_graph_lower(): + raise NotImplementedError( + "a long-range module is evaluated on the NeighborGraph lower, " + "which this descriptor does not provide" + ) + if "energy" not in self.fitting_output_def().keys(): + raise ValueError( + "a long-range module adds to the energy output, which this " + "fitting does not have" + ) + if self.atom_exclude_types: + raise NotImplementedError( + "`atom_exclude_types` is not supported with a long-range module, " + "which attributes each pair energy to both atoms of the pair" + ) + if self.pair_exclude_types and any( + name != ALL_PAIRS for name in self.long_range.required_topologies() + ): + raise NotImplementedError( + "`pair_exclude_types` is not supported with a long-range module " + "that serves periodic systems, whose reciprocal-space sum is " + "global and excludes no pair" + ) + if self.long_range.uses_charge_spin() and not self.add_chg_spin_ebd: + raise ValueError( + "the long-range module reads the total charge from " + "`charge_spin`, but the descriptor is not configured with the " + "charge/spin condition (`add_chg_spin_ebd`)" + ) + def init_vacuum_conditions(self) -> None: """Build the conditioning inputs of the reference atoms for the type map. @@ -214,20 +267,49 @@ def uses_graph_lower(self) -> bool: return bool(self.descriptor.uses_graph_lower()) def fused_decomposition(self) -> "tuple[DPAtomicModel, None]": - """A descriptor-fitting model is its own learned part.""" + """A descriptor-fitting model is its own learned part. + + The learned part carries the long-range module of the model, if any; + a fused pipeline that does not evaluate long-range modules declines a + learned part with one. + """ return self, None + def has_long_range(self) -> bool: + """Whether this model carries a long-range module.""" + return self.long_range is not None + + def required_topologies(self) -> tuple[str, ...]: + """The topologies of the long-range module.""" + if self.long_range is None: + return () + return self.long_range.required_topologies() + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of the long-range module.""" + if self.long_range is None: + return {} + return self.long_range.topology_params() + def has_message_passing_across_ranks(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.has_message_passing_across_ranks()) def supports_edge_parallel(self) -> bool: - """Delegates to this model's own descriptor.""" - return bool(self.descriptor.supports_edge_parallel()) + """Delegates to this model's own descriptor. + + A long-range module sums over complete systems, so a model carrying + one does not run under domain decomposition. + """ + return self.long_range is None and bool( + self.descriptor.supports_edge_parallel() + ) def dense_lower_supports_comm(self) -> bool: - """Delegates to this model's own descriptor.""" - return bool(self.descriptor.dense_lower_supports_comm()) + """Delegates to this model's own descriptor; a long-range module has no dense lower.""" + return self.long_range is None and bool( + self.descriptor.dense_lower_supports_comm() + ) def uses_compact_edge_pairs(self) -> bool: """Delegates to this model's own descriptor.""" @@ -267,6 +349,8 @@ def set_case_embd(self, case_idx: int) -> None: typically concatenated with the output of the descriptor and fed into the fitting net. """ self.fitting_net.set_case_embd(case_idx) + if self.long_range is not None: + self.long_range.set_case_embd(case_idx) def mixed_types(self) -> bool: """If true, the model @@ -298,6 +382,12 @@ def enable_compression( ) -> None: """Call descriptor enable_compression(). + A long-range module keeps its networks: it reads the output of the + compressed descriptor like that of the uncompressed one. A compressed + descriptor folds the charge/spin condition into its tables, so a + constraint that reads the total charge from the condition is folded + to the same state, the default condition of the descriptor. + Parameters ---------- min_nbor_dist @@ -318,6 +408,8 @@ def enable_compression( table_stride_2, check_frequency, ) + if self.long_range is not None and self.long_range.uses_charge_spin(): + self.long_range.fold_charge_state(self.descriptor.get_default_chg_spin()) def forward_atomic( self, @@ -358,6 +450,11 @@ def forward_atomic( the result dict, defined by the `FittingOutputDef`. """ + if self.long_range is not None: + raise NotImplementedError( + "a long-range module is evaluated on the NeighborGraph lower " + "only; the dense (nlist) lower cannot evaluate this model" + ) nframes, nloc, nnei = nlist.shape atype = xp_take_first_n(extended_atype, 1, nloc) @@ -413,11 +510,11 @@ def vacuum_descriptor(self) -> Array: conditioned as the neutral ground-state atom: zero charge with the ground-state multiplicity when the descriptor takes the charge/spin condition, and a spin vector of one Bohr magneton per unpaired - electron when it takes the native spin. A descriptor with a graph - lower is evaluated on the neighbor graph of those single-atom frames, - which holds no real edge; a dense descriptor on a neighbor list of its - own capacity filled with padding, since its equations are sized by - that capacity. + electron when it takes the native spin. + + Descriptors that opt into graph vacuum evaluation use isolated nodes, + independently of their advertised neighbor capacity; other descriptors + retain their configured all-masked dense neighbor list. Returns ------- @@ -429,39 +526,36 @@ def vacuum_descriptor(self) -> Array: xp = array_api_compat.array_namespace(bias) device = array_api_compat.device(bias) ntypes = self.get_ntypes() - coord = xp.zeros((ntypes, 1, 3), dtype=bias.dtype, device=device) - atype = xp.reshape( - xp.arange(ntypes, dtype=xp.int64, device=device), (ntypes, 1) - ) + atype = xp.arange(ntypes, dtype=xp.int64, device=device) conditions = { name: xp.asarray(table, dtype=bias.dtype, device=device) for name, table in self.vacuum_conditions().items() } - if self.descriptor.uses_graph_lower(): + if self.uses_graph_lower() and self.descriptor.supports_graph_vacuum(): from deepmd.dpmodel.utils.neighbor_graph import ( - build_neighbor_graph, + NeighborGraph, ) - graph = build_neighbor_graph( - coord, atype, None, self.descriptor.get_rcut(), with_csr=True - ) - # Only descriptors whose ``call_graph`` declares a condition take - # the keyword; a declared condition the model does not use is None. - spin_kwargs = ( - {"spin": conditions.get("spin")} if self.supports_native_spin() else {} - ) - charge_spin_kwargs = ( - {"charge_spin": conditions.get("charge_spin")} - if self.supports_charge_spin - else {} + # One isolated node per frame, with the same two masked guard + # edges as the graph builders. Never size this graph from sel: + # graph-only descriptors may report an unbounded sentinel there. + # The graph interface also carries native spin, which the dense + # compatibility adapters need not support. + graph = NeighborGraph( + n_node=xp.ones((ntypes,), dtype=xp.int64, device=device), + edge_index=xp.zeros((2, 2), dtype=xp.int64, device=device), + edge_vec=xp.zeros((2, 3), dtype=bias.dtype, device=device), + edge_mask=xp.zeros((2,), dtype=xp.bool, device=device), ) return self.descriptor.call_graph( graph, - xp.reshape(atype, (ntypes,)), + atype, type_embedding=self.descriptor.graph_type_embedding_table(), - **spin_kwargs, - **charge_spin_kwargs, + **conditions, )[0] + + coord = xp.zeros((ntypes, 1, 3), dtype=bias.dtype, device=device) + atype = xp.reshape(atype, (ntypes, 1)) nlist = xp.full( (ntypes, 1, self.descriptor.get_nnei()), -1, @@ -575,13 +669,16 @@ def forward_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict[str, Array]: """Graph analogue of :meth:`forward_atomic` on the flat node axis. Runs the descriptor ``call_graph`` then the fitting ``call_graph`` PER NODE and returns the raw fitting dict on the flat ``(N, *)`` axis (no reduction or masking; the wrapper handles those). ``fparam`` is gathered to nodes by - ``frame_id`` so each node sees its frame's parameter. + ``frame_id`` so each node sees its frame's parameter. The atomic energy + of the long-range module is added to the energy output of the fitting. Parameters ---------- @@ -607,6 +704,11 @@ def forward_atomic_graph( ``call_graph`` (the message-passing part). ``None`` for non-parallel inference (default). Mirrors :meth:`forward_atomic`'s ``comm_dict`` on the dense route. + topologies + the topologies of :meth:`required_topologies` by name. + node_mask + flat mask of the real and owned nodes, which carry the fields of + the long-range module. Returns ------- @@ -638,11 +740,28 @@ def forward_atomic_graph( graph.n_local, atype.shape[0], ) + # The long-range module reads the frame conditions of the real frames, + # with the defaults the descriptor and fitting fall back on. + n_node = graph.n_node + if self.long_range is not None: + # The constraint and a compressed descriptor serve one charge + # state; compressing the atomic model folds both to it. + if self.long_range.uses_charge_spin() and self.get_dim_chg_spin() == 0: + raise ValueError( + "the descriptor folds the charge/spin condition into its " + "compressed tables, while the long-range charge constraint " + "still reads it; compress the atomic model, which folds " + "both to one state" + ) + fparam, charge_spin_frame = self._frame_conditions( + fparam, charge_spin, n_node.shape[0], graph.edge_vec + ) # === Step 2. Vacuum reference nodes === - # The reference atoms are appended as single-node frames, so the same - # descriptor call yields the vacuum descriptor of every type. + # Opted-in descriptors evaluate the reference in the same graph call. + # Legacy descriptors must retain the dense reference used by folding. vacuum_ref = self.fitting_net.needs_vacuum_descriptor() - if vacuum_ref: + graph_vacuum = vacuum_ref and self.descriptor.supports_graph_vacuum() + if graph_vacuum: graph, atype, charge_spin, spin = self.append_vacuum_frames( graph, atype, charge_spin, spin ) @@ -671,19 +790,21 @@ def forward_atomic_graph( gg, rot_mat = outputs[0], outputs[1] node_gate = outputs[2] if len(outputs) > 2 else None vacuum_kwargs = {} - if vacuum_ref: + if graph_vacuum: vacuum_kwargs["vacuum_descriptor"] = gg[n_real:] gg = gg[:n_real] atype = atype[:n_real] rot_mat = None if rot_mat is None else rot_mat[:n_real] node_gate = None if node_gate is None else node_gate[:n_real] + elif vacuum_ref: + vacuum_kwargs["vacuum_descriptor"] = self.vacuum_descriptor() # === Step 4. Fitting === # The source gate of a bridged descriptor fades the learned part of # the fitting output: an atom of a frozen pair keeps only the bias of # its type, so along the closing of a pair the learned energy follows # the switch linearly while the analytical pair term takes over. gate_kwargs = {} if node_gate is None else {"node_gate": node_gate} - return self.fitting_net.call_graph( + ret = self.fitting_net.call_graph( gg, atype, gr=rot_mat, @@ -694,6 +815,71 @@ def forward_atomic_graph( **vacuum_kwargs, **gate_kwargs, ) + # === Step 5. Long-range module === + # The module reads the descriptor of the real nodes; the reference + # nodes appended above take no part in it. + if self.long_range is not None: + if topologies is None or node_mask is None: + raise ValueError( + "the long-range module of this model reads the topologies " + f"{self.required_topologies()} and the node mask, which were " + "not given" + ) + ret["energy"] = ret["energy"] + self.long_range.call_graph( + gg, + atype, + node_mask, + n_node, + topologies, + fparam=fparam, + aparam=aparam_node, + charge_spin=charge_spin_frame, + ) + return ret + + def _frame_conditions( + self, + fparam: Array | None, + charge_spin: Array | None, + nf: int, + reference: Array, + ) -> tuple[Array | None, Array | None]: + """Frame parameters and charge/spin condition with their defaults resolved. + + Parameters + ---------- + fparam + Frame parameters with shape (nf, numb_fparam), or None. + charge_spin + Charge/spin condition with shape (nf, 2) or (1, 2), or None. + nf + The number of frames. + reference + An array giving the namespace, dtype and device of the defaults. + + Returns + ------- + tuple + The frame parameters with shape (nf, numb_fparam) and the + charge/spin condition with shape (nf, 2), each None when the model + takes no such input or it has no default. + """ + xp = array_api_compat.array_namespace(reference) + device = array_api_compat.device(reference) + if fparam is None and self.has_default_fparam(): + default_fparam = xp.asarray( + self.get_default_fparam(), dtype=reference.dtype, device=device + ) + fparam = xp.tile(xp.reshape(default_fparam, (1, -1)), (nf, 1)) + if not self.add_chg_spin_ebd: + return fparam, None + if charge_spin is None: + default_cs = self.descriptor.get_default_chg_spin() + if default_cs is None: + return fparam, None + charge_spin = xp.asarray(default_cs, dtype=reference.dtype, device=device) + charge_spin = xp.broadcast_to(xp.reshape(charge_spin, (-1, 2)), (nf, 2)) + return fparam, charge_spin def compute_or_load_stat( self, @@ -724,6 +910,10 @@ def compute_or_load_stat( self.fitting_net.compute_input_stats( wrapped_sampler, stat_file_path=stat_file_path ) + if self.long_range is not None: + self.long_range.compute_input_stats( + wrapped_sampler, stat_file_path=stat_file_path + ) self._collect_and_set_observed_type( wrapped_sampler, stat_file_path, preset_observed_type ) @@ -747,6 +937,8 @@ def change_type_map( else None, ) self.fitting_net.change_type_map(type_map=type_map) + if self.long_range is not None: + self.long_range.change_type_map(type_map=type_map) self.init_vacuum_conditions() def compute_fitting_input_stat( @@ -756,6 +948,9 @@ def compute_fitting_input_stat( ) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + The networks of the long-range module take the inputs of the fitting + and receive their statistics as well. + Parameters ---------- sample_merged : Union[Callable[[], list[dict]], list[dict]] @@ -772,6 +967,37 @@ def compute_fitting_input_stat( protection=self.data_stat_protect, stat_file_path=stat_file_path, ) + if self.long_range is not None: + self.long_range.compute_input_stats( + sample_merged, + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) + + def compute_long_range_input_stat( + self, + sampled_func: Callable[[], list[dict]], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics of the long-range module alone. + + A fine-tuning run that adds a long-range module to a pretrained model + calls it for the new module, whose statistics, unlike those of the + inherited descriptor and fitting, the pretrained model cannot provide. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data + systems. + stat_file_path + The path to the stat file. + """ + self.long_range.compute_input_stats( + self._make_wrapped_sampler(sampled_func), + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) def serialize(self) -> dict: dd = super().serialize() @@ -785,6 +1011,12 @@ def serialize(self) -> dict: "fitting": self.fitting_net.serialize(), } ) + # A record with a long-range module is version 3, which a reader that + # cannot evaluate the module rejects; every other record keeps + # version 2. + if self.long_range is not None: + dd["@version"] = 3 + dd["long_range"] = self.long_range.serialize() return dd # for subclass overridden @@ -792,17 +1024,27 @@ def serialize(self) -> dict: """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls: type | None = BaseLongRange + """The base long-range module class; ``None`` for a backend without one.""" @classmethod def deserialize(cls, data: dict[str, Any]) -> "DPAtomicModel": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 2, 2) + check_version_compatibility(data.pop("@version", 1), 3, 2) data.pop("@class") data.pop("type") descriptor_obj = cls.base_descriptor_cls.deserialize(data.pop("descriptor")) fitting_obj = cls.base_fitting_cls.deserialize(data.pop("fitting")) data["descriptor"] = descriptor_obj data["fitting"] = fitting_obj + long_range = data.pop("long_range", None) + if long_range is not None: + if cls.base_long_range_cls is None: + raise NotImplementedError( + "this backend cannot evaluate the long-range module of the " + "model record" + ) + data["long_range"] = cls.base_long_range_cls.deserialize(long_range) obj = super().deserialize(data) return obj diff --git a/deepmd/dpmodel/atomic_model/inner_potential.py b/deepmd/dpmodel/atomic_model/inner_potential.py index 4d3c51f2d7..749e65ee14 100644 --- a/deepmd/dpmodel/atomic_model/inner_potential.py +++ b/deepmd/dpmodel/atomic_model/inner_potential.py @@ -412,7 +412,8 @@ def call( r = xp.clip(r, min=1e-10) # Padded atoms carry type -1; their edges are masked below, so the # table lookup only needs a valid index for them. - atype_row = xp.clip(xp.astype(atype_flat, xp.int64), min=0) + atype_row = xp.astype(atype_flat, xp.int64) + atype_row = xp.maximum(atype_row, xp.zeros_like(atype_row)) # Row ``center * (T + 1) + neighbor``, the ordered-pair index the # fused kernels use, with the destination as the center. row = xp.take( @@ -606,12 +607,15 @@ def forward_atomic_graph( charge_spin: Any = None, spin: Any = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Any = None, ) -> dict: """Evaluate the analytical per-atom energy on the flat node axis. - ``fparam``/``aparam``/``charge_spin``/``spin``/``comm_dict`` are - accepted for pipeline-signature compatibility and ignored (the - analytical term conditions on geometry and types only). + ``fparam``/``aparam``/``charge_spin``/``spin``/``comm_dict``/ + ``topologies``/``node_mask`` are accepted for pipeline-signature + compatibility and ignored (the analytical term conditions on + geometry and types only; the caller masks its outputs). Parameters ---------- diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index e224cbdc2b..0d692fe0a5 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -12,6 +12,9 @@ from deepmd.dpmodel.array_api import ( Array, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, +) from deepmd.dpmodel.utils.nlist import ( build_multiple_neighbor_list, get_multiple_nlist_key, @@ -34,6 +37,9 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from .inner_potential import ( + InnerPotentialAtomicModel, +) from .pairtab_atomic_model import ( PairTabAtomicModel, ) @@ -112,6 +118,47 @@ def __init__( + ", ".join(f"{type(m).__name__}={dim_of(m)}" for m in models) ) + # A long-range module is evaluated on the NeighborGraph lower and + # attributes each pair energy to both atoms of the pair. It is not + # combined with a sub-model without that lower, which sends the whole + # composition to the dense route, with a bridging term, which defines + # the close-pair limit of a learned local energy, nor with atom + # exclusion, which removes atomic outputs. + long_range = any(m.has_long_range() for m in models) + if long_range and not all(m.uses_graph_lower() for m in models): + raise NotImplementedError( + "a sub-model carrying a long-range module is evaluated on the " + "NeighborGraph lower, which every sub-model of the composition " + "must provide" + ) + if long_range and any(isinstance(m, InnerPotentialAtomicModel) for m in models): + raise NotImplementedError( + "an inner_potential bridging term cannot be combined with a " + "sub-model carrying a long-range module" + ) + if self.atom_exclude_types and long_range: + raise NotImplementedError( + "`atom_exclude_types` is not supported with a sub-model carrying " + "a long-range module" + ) + # The composition builds every topology once for all sub-models: the + # pairs of a periodic batch are guards only, which would silence a + # module serving non-periodic systems alone, and a reciprocal-space + # topology is built with the parameters of a single module. + declared = [m.required_topologies() for m in models if m.has_long_range()] + periodic = any(name != ALL_PAIRS for names in declared for name in names) + if len(declared) > 1 and periodic: + raise NotImplementedError( + "a sub-model whose long-range module serves periodic systems " + "takes no other sub-model carrying a long-range module beside it" + ) + if self.pair_exclude_types and periodic: + raise NotImplementedError( + "`pair_exclude_types` is not supported with a sub-model whose " + "long-range module serves periodic systems, whose " + "reciprocal-space sum is global and excludes no pair" + ) + self.models = models self.type_map = type_map self._rebuild_mapping_state() @@ -377,6 +424,8 @@ def compute_fitting_input_stats( for model in self.models: if isinstance(model, DPAtomicModel): model.fitting_net.compute_input_stats(sampled_func) + if model.long_range is not None: + model.long_range.compute_input_stats(sampled_func) def supports_native_spin(self) -> bool: """Spin-capable when ANY child consumes the spin input. @@ -389,6 +438,30 @@ def supports_native_spin(self) -> bool: """ return any(m.supports_native_spin() for m in self.models) + def has_long_range(self) -> bool: + """ANY child carrying a long-range module makes the composition carry one.""" + return any(m.has_long_range() for m in self.models) + + def required_topologies(self) -> tuple[str, ...]: + """Every topology a child reads, each named once.""" + return tuple( + dict.fromkeys( + name for model in self.models for name in model.required_topologies() + ) + ) + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of the sub-models. + + At most one sub-model declares a reciprocal-space topology, the only + kind with build parameters, so that the parameters of every topology + come from one sub-model. + """ + params: dict[str, dict[str, Any]] = {} + for model in self.models: + params.update(model.topology_params()) + return params + def forward_atomic_graph( self, graph: Any, @@ -398,14 +471,18 @@ def forward_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict[str, Array]: """Graph-route linear combination on the flat node axis. - Every child consumes the SAME graph, so on autograd backends the - shared ``graph.edge_vec`` leaf makes the summed energy's force and - virial exactly the sum of the children's -- one edge backward - covers the whole composition (this is what makes analytical - bridging terms compose with the learned model for free). + Every child consumes the SAME graph and topologies, so on autograd + backends the shared geometry leaves make the summed energy's force + and virial exactly the sum of the children's -- one backward covers + the whole composition (this is what makes analytical bridging terms + compose with the learned model for free). Each child evaluates its + complete energy, including its long-range module, before its weight + applies. Only constant weights are supported here (``"sum"``/``"mean"`` or a per-child float list); the distance-switched ZBL-interpolation @@ -429,6 +506,12 @@ def forward_atomic_graph( flat (N, 3) per-node spin, forwarded to every child. comm_dict MPI communication metadata, forwarded to every child. + topologies + the topologies of :meth:`required_topologies`, forwarded to every + child. + node_mask + flat mask of the nodes whose outputs count, forwarded to every + child, which combines it with its own exclusions. Returns ------- @@ -472,6 +555,8 @@ def forward_atomic_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, + node_mask=node_mask, ) contrib = ret["energy"] * ww energy = contrib if energy is None else energy + contrib @@ -933,7 +1018,7 @@ def _compute_weight( # use the larger rr based on nlist nlist_larger = zbl_nlist if zbl_nnei >= dp_nnei else dp_nlist - masked_nlist = xp.clip(nlist_larger, 0, None) + masked_nlist = xp.maximum(nlist_larger, xp.zeros_like(nlist_larger)) pairwise_rr = PairTabAtomicModel._get_pairwise_dist( extended_coord, masked_nlist ) diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 41a1741e14..2d16c5559a 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -434,7 +434,12 @@ def _extract_spline_coefficient( expanded_idx = xp.broadcast_to( idx[..., xp.newaxis, xp.newaxis], (*idx.shape, 1, 4) ) - clipped_indices = xp.astype(xp.clip(expanded_idx, 0, nspline - 1), xp.int64) + clipped_indices = xp.astype( + xp.minimum( + xp.maximum(expanded_idx, xp.zeros_like(expanded_idx)), nspline - 1 + ), + xp.int64, + ) # (nframes, nloc, nnei, 4) final_coef = xp.squeeze( diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 7e264fb3e9..89e47669e4 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -2512,6 +2512,10 @@ def dense_lower_supports_comm(self) -> bool: """ return False + def supports_graph_vacuum(self) -> bool: + """The dense adapter uses the same graph core; native spin needs it.""" + return True + def uses_graph_lower(self) -> bool: """Whether this descriptor supports the sel-free NeighborGraph lower. diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index 7dbd9c7f3d..d43c9255bd 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -90,7 +90,6 @@ OrderedPairFiLM, SpinChannels, build_angular_basis, - build_moment_indices, canonicalize_charge_spin, degree_offsets, derive_bispectrum_ranks, @@ -245,6 +244,8 @@ class DescrptDPA4C(NativeOP, BaseDescriptor): _STAT_FRAMES_PER_SAMPLE = 64 _COMPRESSION_BUFFER_NAMES = ( "data", + "basis_table", + "basis_matrix", "info", "pair_film", "pair_mixing", @@ -258,6 +259,10 @@ class DescrptDPA4C(NativeOP, BaseDescriptor): "output_mean", "output_inv_std", ) + # The low-rank radial table and its expansion, which a payload compressed + # without that table does not carry; empty artifacts select the radial + # table itself, so such a payload loads as the snapshot it was. + _OPTIONAL_COMPRESSION_BUFFER_NAMES = ("basis_table", "basis_matrix") def __init__( self, @@ -461,12 +466,9 @@ def __init__( # === Step 4. Lay out the flat moment payload === # Degree zero owns the leading `channels` entries of the flat layout, - # so the non-scalar block is exactly its complement and needs no - # separate degree index. The spin families, when present, are appended - # after the geometric degrees, leaving every geometric offset intact. - channel_index, harmonic_index = build_moment_indices(self.degree_channels) - self.angular_channel_index = channel_index[self.channels :] - self.angular_harmonic_index = harmonic_index[self.channels :] + # and every higher degree follows in its own harmonic-major block. The + # spin families, when present, are appended after the geometric + # degrees, leaving every geometric offset intact. self.degree_offsets = degree_offsets(self.degree_channels) # === Step 5. Initialize the output calibration state === @@ -1146,30 +1148,30 @@ def aggregate_moments( ) xp = array_api_compat.array_namespace(amplitude) - device = array_api_compat.device(amplitude) - channel_index = xp_asarray_nodetach( - xp, - self.angular_channel_index, - device=device, - ) - harmonic_index = xp_asarray_nodetach( - xp, - self.angular_harmonic_index, - device=device, - ) + n_edge = amplitude.shape[0] # Payload layout: [chi^2, chi^4, degree zero (C_0), degrees one and # above (S - C_0)]. Degree zero is the whole edge amplitude under its - # single envelope, while the non-scalar block gathers an amplitude and - # a harmonic per flat moment coordinate and carries a second envelope. + # single envelope. Degree l is the outer product of its 2l + 1 + # harmonics with the leading C_l channels of the amplitude, laid out + # harmonic-major and carrying a second envelope. The products are + # formed by broadcasting, so their backward reduces over the + # broadcast axes. + angular = [ + xp.reshape( + basis[:, degree * degree : (degree + 1) * (degree + 1), None] + * amplitude[:, None, :width], + (n_edge, (2 * degree + 1) * width), + ) + for degree, width in enumerate(self.degree_channels) + if degree > 0 + ] envelope_squared = envelope * envelope parts = [ envelope_squared[:, None], (envelope_squared * envelope_squared)[:, None], amplitude, - xp.take(amplitude, channel_index, axis=1) - * xp.take(basis, harmonic_index, axis=1) - * envelope[:, None], + xp.concat(angular, axis=1) * envelope[:, None], ] if spin_payload is not None: parts.append(spin_payload) @@ -1931,11 +1933,36 @@ def deserialize(cls, data: dict) -> DescrptDPA4C: obj._load_compression(compression) return obj + @classmethod + def _compression_variables(cls, compression: dict) -> dict: + """Return every compressed-inference artifact of a payload. + + Parameters + ---------- + compression + The ``compress`` entry of a serialized descriptor. + + Returns + ------- + dict + The artifact of every name in ``_COMPRESSION_BUFFER_NAMES``, an + empty float32 array for an optional one the payload omits. + """ + variables = compression["@variables"] + return { + name: ( + np.empty(0, dtype=np.float32) + if name in cls._OPTIONAL_COMPRESSION_BUFFER_NAMES + and name not in variables + else variables[name] + ) + for name in cls._COMPRESSION_BUFFER_NAMES + } + def _load_compression(self, compression: dict) -> None: """Restore the compressed-inference artifact payload.""" - variables = compression["@variables"] - for name in self._COMPRESSION_BUFFER_NAMES: - setattr(self, f"compress_{name}", variables[name]) + for name, value in self._compression_variables(compression).items(): + setattr(self, f"compress_{name}", value) self.compress = True @classmethod @@ -2118,6 +2145,10 @@ def get_env_protection(self) -> float: """Return the direction regularization scale in Å.""" return self._EPS + def supports_graph_vacuum(self) -> bool: + """Isolated references use the graph core without a dense sel capacity.""" + return True + def uses_graph_lower(self) -> bool: """Return whether graph-native lowering is supported.""" return True diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py index 460c7970e6..8bd43d4cc0 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py @@ -14,7 +14,6 @@ from .geometry import ( MAX_ANGULAR_DEGREE, build_angular_basis, - build_moment_indices, degree_offsets, derive_degree_channels, packed_l2_to_stf, @@ -41,7 +40,6 @@ "SpinChannels", "build_angular_basis", "build_bispectrum_layout", - "build_moment_indices", "canonicalize_charge_spin", "degree_offsets", "derive_bispectrum_ranks", diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py index 766a9d060c..a9378f4ea1 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py @@ -11,7 +11,6 @@ ) import array_api_compat -import numpy as np if TYPE_CHECKING: from collections.abc import ( @@ -236,36 +235,3 @@ def degree_offsets(degree_channels: Sequence[int]) -> tuple[int, ...]: for degree, width in enumerate(degree_channels): offsets.append(offsets[-1] + (2 * degree + 1) * int(width)) return tuple(offsets) - - -def build_moment_indices( - degree_channels: Sequence[int], -) -> tuple[np.ndarray, np.ndarray]: - """Build channel and harmonic indices for one flat edge payload. - - Every degree reads the leading channels of the shared radial map, so the - channel index of a degree is simply ``range(degree_channels[degree])``. - - Parameters - ---------- - degree_channels - Number of channels for degrees zero through ``lmax``. - - Returns - ------- - channel_index - Edge-amplitude channel indices with shape ``(S,)``. - harmonic_index - Packed harmonic indices with shape ``(S,)``, where - ``S = sum((2 * l + 1) * degree_channels[l])``. - """ - channel_index = [] - harmonic_index = [] - for degree, width in enumerate(degree_channels): - for component in range(2 * degree + 1): - channel_index.extend(range(int(width))) - harmonic_index.extend([degree * degree + component] * int(width)) - return ( - np.asarray(channel_index, dtype=np.int64), - np.asarray(harmonic_index, dtype=np.int64), - ) diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index b0598e9735..9ea052c368 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -287,6 +287,18 @@ def graph_type_embedding_table(self) -> Any | None: """ return None + def supports_graph_vacuum(self) -> bool: + """Whether isolated-atom references can use the graph interface. + + Opting in requires ``call_graph`` to produce the same isolated + reference as the dense compatibility path, when one exists; + graph-only descriptors may also opt in. This is distinct from + ``uses_graph_lower``: some legacy descriptors retain different + padding semantics on their dense and graph paths, so switching + their reference would change the model's energy zero. + """ + return False + @abstractmethod def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 37b0a1b731..126b51e721 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -14,6 +14,9 @@ from .invar_fitting import ( InvarFitting, ) +from .latent_charge import ( + LatentChargeFitting, +) from .make_base_fitting import ( make_base_fitting, ) @@ -29,6 +32,7 @@ "DipoleFitting", "EnergyFittingNet", "InvarFitting", + "LatentChargeFitting", "PolarFitting", "PropertyFittingNet", "SeZMEnergyFittingNet", diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 3693cc1c3c..f6f22aa493 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -944,6 +944,49 @@ def readout_reference(self) -> Array: device=array_api_compat.device(self.bias_atom_e), ) + def fold_case_embedding(self) -> None: + """Fold the case embedding into the first layer so the call needs no case columns. + + The case embedding is the same trailing block of input columns for + every atom, so its product with the rows of the first-layer weight + that read it is a constant of the pre-activation. Adding that product + to the bias and removing the rows yields a fitting without the case + embedding that gives the same outputs. Export folds the fixed case of + a deployed model, so that the frozen branch of a multi-task model with + a shared fitting network is a plain fitting network. The fold rebuilds + the first layers in place, so it applies to a fitting that no other + branch shares. A first layer whose residual adds its input to the + output also carries the case embedding past the activation; such a + fitting keeps it. + """ + if self.dim_case_embd == 0: + return + networks = [ + self.nets[ii] for ii in range(self.ntypes if not self.mixed_types else 1) + ] + for net in networks: + width_in, width_out = net[0].w.shape + if net[0].resnet and width_out in (width_in, 2 * width_in): + return + case_embd = to_numpy_array(self.case_embd).astype(np.float64) + for net in networks: + data = net[0].serialize() + variables = data["@variables"] + weight = variables["w"] + kept = weight.shape[0] - self.dim_case_embd + variables["b"] = ( + variables["b"].astype(np.float64) + + case_embd @ weight[kept:].astype(np.float64) + ).astype(weight.dtype) + variables["w"] = weight[:kept] + # The residual is inactive at the original input width, checked + # above; clearing it keeps the narrower input from activating it. + data["resnet"] = False + net[0] = type(net[0]).deserialize(data) + net.in_dim = kept + self.dim_case_embd = 0 + self.case_embd = None + def _call_common( self, descriptor: Array, diff --git a/deepmd/dpmodel/fitting/latent_charge.py b/deepmd/dpmodel/fitting/latent_charge.py new file mode 100644 index 0000000000..4967879f1f --- /dev/null +++ b/deepmd/dpmodel/fitting/latent_charge.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) + +from .invar_fitting import ( + InvarFitting, +) + + +@fitting_check_output +class LatentChargeFitting(InvarFitting): + r"""Per-atom latent charges predicted from the invariant descriptor. + + The network maps the descriptor of an atom, together with the frame + parameters, atomic parameters and case embedding of the energy fitting it + accompanies, to ``dim_out`` latent charge channels in units of the + elementary charge :math:`e`. The charges are an intermediate field of a + long-range module rather than a model output: the module constrains and + pairs them, and the potential energy is the only quantity they are + trained through. + + The constructor arguments are those of :class:`InvarFitting` with the + output variable fixed to ``latent_charge``. + """ + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def serialize(self) -> dict: + data = super().serialize() + data["type"] = "latent_charge" + return data diff --git a/deepmd/dpmodel/long_range/__init__.py b/deepmd/dpmodel/long_range/__init__.py new file mode 100644 index 0000000000..4a4aad0d7c --- /dev/null +++ b/deepmd/dpmodel/long_range/__init__.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Long-range modules of the atomic energy of a descriptor-fitting model. + +A long-range module is a component of the atomic energy model beside its +energy fitting: it reads the descriptor features shared with that fitting and +the topologies it declares, and contributes an energy per atom. Every consumer +of the atomic model therefore evaluates the complete potential. +""" + +from .base_long_range import ( + BaseLongRange, + check_long_range_support, + make_base_long_range, +) +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, + LESKernel, + SOGKernel, +) +from .latent_charge import ( + LatentChargeLongRange, +) + +__all__ = [ + "COULOMB_CONSTANT", + "BaseInteractionKernel", + "BaseLongRange", + "LESKernel", + "LatentChargeLongRange", + "SOGKernel", + "check_long_range_support", + "make_base_long_range", +] diff --git a/deepmd/dpmodel/long_range/base_long_range.py b/deepmd/dpmodel/long_range/base_long_range.py new file mode 100644 index 0000000000..d00617fc1c --- /dev/null +++ b/deepmd/dpmodel/long_range/base_long_range.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + ABC, + abstractmethod, +) +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.plugin import ( + PluginVariant, + make_plugin_registry, +) + + +def check_long_range_support(descriptor: Any) -> None: + """Admit only the descriptors whose combination with long-range modules is validated. + + A long-range module reads nothing but the invariant node features every + graph descriptor returns, so this gate is a statement of validation, not + of implementation: a further descriptor is admitted by adding it here once + its combination has been tested. + + Parameters + ---------- + descriptor + The descriptor of the atomic model receiving the long-range module. + + Raises + ------ + NotImplementedError + If the descriptor is not among the validated ones. + """ + from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, + ) + + if not isinstance(descriptor, DescrptDPA4C): + raise NotImplementedError( + "long-range modules are currently validated with the DPA4C " + f"descriptor only, got {type(descriptor).__name__}" + ) + + +def make_base_long_range(t_tensor: Any) -> type: + """Make the base class of the long-range modules of one backend. + + Each backend owns its registry, so a configuration builds the long-range + module native to the backend that reads it. + + Parameters + ---------- + t_tensor + The type of the tensor, used in the type hints. + """ + + class BLR(ABC, PluginVariant, make_plugin_registry("long range")): + """Long-range contribution to the atomic energy of a descriptor-fitting model. + + A long-range module reads the invariant node features of the + descriptor and the topologies it declares, and returns an energy per + node in eV. The atomic model adds it to the output of its energy + fitting before the output statistics and masks, so every consumer of + the atomic model sees one complete potential. The module owns the + attribution of its energy to the nodes; the per-frame energy is their + sum. + """ + + @abstractmethod + def required_topologies(self) -> tuple[str, ...]: + """Names of the topologies the module reads beside the neighbor graph.""" + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies, by name; none by default.""" + return {} + + @abstractmethod + def uses_charge_spin(self) -> bool: + """Whether the module reads the frame-level charge/spin condition.""" + + @abstractmethod + def call_graph( + self, + descriptor: t_tensor, + atype: t_tensor, + node_mask: t_tensor, + n_node: t_tensor, + topologies: dict[str, Any], + fparam: t_tensor | None = None, + aparam: t_tensor | None = None, + charge_spin: t_tensor | None = None, + ) -> t_tensor: + """Evaluate the atomic energy of the module on the flat node axis. + + Parameters + ---------- + descriptor + Invariant node features with shape ``(N, dim_descrpt)``. + atype + Node types with shape ``(N,)``; invalid nodes carry any valid + type. + node_mask + Nodes carrying fields with shape ``(N,)``: the real and owned + nodes. + n_node + Per-frame node counts with shape ``(nf,)``. + topologies + The topologies of :meth:`required_topologies` by name. + fparam + Frame parameters with shape ``(nf, numb_fparam)``, defaults + resolved. + aparam + Atomic parameters with shape ``(N, numb_aparam)``. + charge_spin + Frame charge/spin condition with shape ``(nf, 2)``, defaults + resolved, when the model takes one. + + Returns + ------- + t_tensor + The atomic energy with shape ``(N, 1)``, in eV. + """ + + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics of the networks of the module.""" + + def change_type_map(self, type_map: list[str]) -> None: + """Remap the type-dependent state of the module to a new type map.""" + + def set_case_embd(self, case_idx: int) -> None: + """Select the case embedding of the networks of the module.""" + + def fold_charge_state(self, charge_spin: list[float]) -> None: + """Specialize the module to one frame charge/spin condition. + + Called when the descriptor of the model is compressed, which folds + the condition into its tables, for a module whose + :meth:`uses_charge_spin` holds. + + Raises + ------ + NotImplementedError + If the module cannot be specialized. + """ + raise NotImplementedError( + f"{type(self).__name__} cannot be specialized to a charge/spin " + "condition, which the compression of the model requires" + ) + + def charge_state_constants(self) -> list[t_tensor]: + """The tensors that carry the folded charge/spin condition. + + A deployment that serves another condition writes the outputs of + :meth:`rebuild_charge_state` over them, in this order. + """ + return [] + + def rebuild_charge_state(self, charge_spin: t_tensor) -> list[t_tensor]: + """Build the tensors of :meth:`charge_state_constants` for a condition. + + Parameters + ---------- + charge_spin + The condition with shape ``(1, 2)``. + """ + return [] + + @abstractmethod + def serialize(self) -> dict: + """Serialize the module to a dict.""" + + @classmethod + def deserialize(cls, data: dict) -> "BLR": + """Deserialize a module, dispatching on its registered type.""" + if cls is BLR: + return BLR.get_class_by_type(data["type"]).deserialize(data) + raise NotImplementedError(f"Not implemented in class {cls.__name__}") + + return BLR + + +BaseLongRange = make_base_long_range(np.ndarray) diff --git a/deepmd/dpmodel/long_range/kernel.py b/deepmd/dpmodel/long_range/kernel.py new file mode 100644 index 0000000000..911dfba5dd --- /dev/null +++ b/deepmd/dpmodel/long_range/kernel.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Pair interaction kernels of the latent-charge long-range modules. + +A kernel is a function :math:`K(r)` of the distance between two charges, in +Å⁻¹, so that charges in units of :math:`e` interact with the energy +:math:`k_e q_i q_j K(r_{ij})` in eV, where :math:`k_e` is the Coulomb +constant. A kernel owns its parameters; widths are stored through their +logarithm, which keeps them positive however training moves them. +""" + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, + xp_asarray_nodetach, + xp_erf, +) +from deepmd.dpmodel.common import ( + PRECISION_DICT, + NativeOP, + to_numpy_array, +) +from deepmd.utils.plugin import ( + make_plugin_registry, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +COULOMB_CONSTANT = 14.3996454784255 +r"""Coulomb constant :math:`e^2 / (4 \pi \varepsilon_0)` in eV Å.""" + + +class BaseInteractionKernel(NativeOP, make_plugin_registry("interaction kernel")): + """Pair interaction kernel :math:`K(r)` in Å⁻¹ with its own parameters. + + Parameters + ---------- + trainable + Whether the parameters of the kernel are trained. + precision + The precision in which the parameters are stored; the kernel is + evaluated in the precision of its input. + """ + + PARAMETER_NAMES: tuple[str, ...] = () + """Attributes holding the trainable parameters of the kernel.""" + + def __init__(self, trainable: bool = True, precision: str = "default") -> None: + self.trainable = bool(trainable) + self.precision = precision + self.prec = PRECISION_DICT[precision.lower()] + + def call(self, r_sq: Array) -> Array: + """Evaluate the kernel on squared pair distances. + + Parameters + ---------- + r_sq + Squared distances with shape ``(P,)``, in Ų. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + return self.evaluate(r_sq, self.coefficients(r_sq)) + + def coefficients(self, like: Array) -> dict[str, Array]: + """The coefficients of the kernel formula, derived from its parameters. + + Parameters + ---------- + like + An array giving the namespace, dtype and device of the coefficients. + + Returns + ------- + dict[str, Array] + Every coefficient with its own shape, differentiable with respect + to the parameters. + """ + raise NotImplementedError + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + """Evaluate the kernel formula with given coefficients. + + Parameters + ---------- + r_sq + Squared distances with shape ``(P,)``, in Ų. + coefficients + The coefficients of :meth:`coefficients`, each with its own shape + or with a leading pair axis of length ``P``. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + raise NotImplementedError + + def kspace(self, k_sq: Array) -> Array: + r"""Fourier transform of the kernel on squared wave numbers. + + :math:`\hat K(k) = \int K(r) e^{-i \mathbf{k} \cdot \mathbf{r}} d^3r`, + so that the energy of charges in a periodic cell of volume ``V`` reads + :math:`(k_e / 2V) \sum_{\mathbf{k} \neq 0} \hat K(k) |\hat\rho(\mathbf{k})|^2`. + + Parameters + ---------- + k_sq + Squared wave numbers with shape ``(nk,)`` or ``(rows, nk)``, in + Å⁻²; every entry must be positive. + + Returns + ------- + Array + The transformed kernel with the shape of ``k_sq``, in ų. + """ + raise NotImplementedError + + def contact_value(self, like: Array) -> Array: + r"""The kernel at contact, :math:`K(0)`, in Å⁻¹. + + The k-space lattice sum of a periodic system contains the self terms, + which the evaluation subtracts as :math:`\frac{1}{2} K(0) \sum_i q_i^2`. + + Parameters + ---------- + like + An array giving the namespace, dtype and device of the result. + """ + raise NotImplementedError + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + """Wave number at which the kernel weight falls below ``tolerance``. + + The Gaussian tail of the Fourier-transformed kernel makes the k-space + truncation exponentially convergent; the k grid of a periodic + evaluation ends at this wave number. + + Parameters + ---------- + tolerance + The relative weight of the slowest-decaying mode at the cutoff. + """ + raise NotImplementedError + + def _parameter(self, name: str, like: Array) -> Array: + """A parameter in the namespace, dtype and device of ``like``, autograd kept.""" + xp = array_api_compat.array_namespace(like) + return xp_asarray_nodetach( + xp, + getattr(self, name), + dtype=like.dtype, + device=array_api_compat.device(like), + ) + + @staticmethod + def _row_values(value: Array, k_sq: Array) -> Array: + """A scalar broadcast to shape ``(..., 1)`` over the leading axes of ``k_sq``. + + The gradient of a parameter of the transform sums over every wave + number. Broadcast over the rows, it sums along every row first, a + reduction a compiled graph spreads over the rows, and then over the + rows, rather than over every wave number in one block. + """ + xp = array_api_compat.array_namespace(k_sq) + return xp.broadcast_to(value, (*k_sq.shape[:-1], 1)) + + def serialize(self) -> dict: + return { + "@class": "InteractionKernel", + "type": self.TYPE, + "@version": 1, + "trainable": self.trainable, + "precision": self.precision, + "@variables": { + name: to_numpy_array(getattr(self, name)) + for name in self.PARAMETER_NAMES + }, + } + + @classmethod + def deserialize(cls, data: dict) -> "BaseInteractionKernel": + if cls is BaseInteractionKernel: + return cls.get_class_by_type(data["type"]).deserialize(data) + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + data.pop("type") + variables = data.pop("@variables") + obj = cls(**data, **cls._shape_kwargs(variables)) + for name in cls.PARAMETER_NAMES: + setattr(obj, name, np.asarray(variables[name], dtype=obj.prec)) + return obj + + @classmethod + def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: + """Constructor arguments that size the parameters stored in ``variables``.""" + return {} + + +# Taylor coefficients of erf(x) / x in powers of x^2, 2 / sqrt(pi) (-1)^n / +# (n! (2n + 1)), and the square of x below which LESKernel evaluates them: there +# the nine terms reach the float64 rounding of erf(x) / x. Above it the closed +# form keeps the relative error of its radial derivative within 2e-14 in +# float64 and 1e-5 in float32. +_ERF_OVER_X_SERIES = ( + 1.1283791670955126, + -0.37612638903183754, + 0.11283791670955126, + -0.026866170645131252, + 0.005223977625442188, + -0.0008548327023450852, + 0.00012055332981789664, + -1.492565035840625e-05, + 1.6462114365889246e-06, +) +_ERF_SERIES_X_SQ = 0.0625 + + +@BaseInteractionKernel.register("les") +class LESKernel(BaseInteractionKernel): + r"""Latent Ewald summation kernel :math:`K(r) = \mathrm{erf}(\alpha r) / r`. + + The kernel is the interaction of two Gaussian charge clouds of width + :math:`1/(\sqrt{2}\alpha)`: it reaches :math:`1/r` beyond a few + :math:`1/\alpha` and stays finite, :math:`2\alpha/\sqrt{\pi}`, at contact. + + Parameters + ---------- + alpha + The inverse width :math:`\alpha > 0`, in Å⁻¹. + trainable + Whether :math:`\alpha` is trained. + precision + The precision in which :math:`\log \alpha` is stored. + """ + + TYPE = "les" + PARAMETER_NAMES = ("log_alpha",) + + def __init__( + self, alpha: float = 1.0, trainable: bool = True, precision: str = "default" + ) -> None: + super().__init__(trainable, precision) + if not np.isfinite(alpha) or alpha <= 0.0: + raise ValueError(f"`alpha` must be finite and positive, got {alpha}") + self.log_alpha = np.array(np.log(alpha), dtype=self.prec) + + def coefficients(self, like: Array) -> dict[str, Array]: + xp = array_api_compat.array_namespace(like) + return {"alpha": xp.exp(self._parameter("log_alpha", like))} + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + xp = array_api_compat.array_namespace(r_sq) + alpha = coefficients["alpha"] + x_sq = alpha**2 * r_sq + # K = alpha erf(x) / x with x = alpha r. The derivative of the closed + # form subtracts two terms that agree to second order in x, so near + # contact erf(x) / x comes from its Taylor series in x^2. Both are + # evaluated on every pair, which keeps the shapes static; each reads a + # value of its own domain where the other is selected, which keeps the + # derivatives of both finite. + near = x_sq < _ERF_SERIES_X_SQ + x_far = xp.sqrt(xp.where(near, xp.ones_like(x_sq), x_sq)) + x_sq_near = xp.where(near, x_sq, xp.zeros_like(x_sq)) + series = xp.full_like(x_sq, _ERF_OVER_X_SERIES[-1]) + for coefficient in reversed(_ERF_OVER_X_SERIES[:-1]): + series = series * x_sq_near + coefficient + return alpha * xp.where(near, series, xp_erf(x_far) / x_far) + + def kspace(self, k_sq: Array) -> Array: + xp = array_api_compat.array_namespace(k_sq) + # erf(alpha r)/r is the interaction of two Gaussian charge clouds of + # width 1/(sqrt(2) alpha); its transform is the damped Coulomb + # 4 pi exp(-k^2 / 4 alpha^2) / k^2. + alpha = self._row_values(xp.exp(self._parameter("log_alpha", k_sq)), k_sq) + return 4.0 * np.pi * xp.exp(-k_sq / (4.0 * alpha * alpha)) / k_sq + + def contact_value(self, like: Array) -> Array: + xp = array_api_compat.array_namespace(like) + alpha = xp.exp(self._parameter("log_alpha", like)) + return 2.0 * alpha / np.sqrt(np.pi) + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + # exp(-k^2 / 4 alpha^2) < tolerance + alpha = float(np.exp(to_numpy_array(self.log_alpha).astype(np.float64))) + return 2.0 * alpha * np.sqrt(np.log(1.0 / tolerance)) + + +@BaseInteractionKernel.register("sog") +class SOGKernel(BaseInteractionKernel): + r"""Sum-of-Gaussians kernel :math:`K(r) = \sum_k a_k \exp(-r^2 / 2\sigma_k^2)`. + + Parameters + ---------- + amplitude + The amplitudes :math:`a_k`, in Å⁻¹, of either sign: a scalar shared by + every Gaussian or one value per Gaussian. Required with ``width``; + with the default widths it defaults to + :math:`a_k = \sqrt{2/\pi} \ln b / \sigma_k`, the trapezoidal + discretization in :math:`\ln s` of + :math:`1/r = \sqrt{2/\pi} \int_0^\infty e^{-r^2/2s^2} s^{-2} \, ds`. + The default kernel then follows :math:`1/r`, less a nearly constant + offset, from a few :math:`\sigma` to a fraction of + :math:`\sigma b^{M-1}`, and stays finite at contact. The offset + exerts no force between fixed charges; through charges that respond + to the positions it does. + width + The widths :math:`\sigma_k > 0`, in Å. Defaults to the geometric + sequence :math:`\sigma_k = \sigma b^k`, :math:`k = 0, \dots, M - 1`. + n_gaussian + The number :math:`M` of Gaussians of the default widths. + base_width + The width :math:`\sigma` of the first default Gaussian, in Å. + width_ratio + The ratio :math:`b > 1` of consecutive default widths. + trainable + Whether the amplitudes and widths are trained. + precision + The precision in which the amplitudes and log-widths are stored. + """ + + TYPE = "sog" + PARAMETER_NAMES = ("amplitude", "log_width") + + def __init__( + self, + amplitude: float | list[float] | None = None, + width: list[float] | None = None, + n_gaussian: int = 12, + base_width: float = 2.180230445405648, + width_ratio: float = 1.62976708826776469, + trainable: bool = True, + precision: str = "default", + ) -> None: + super().__init__(trainable, precision) + if width is None: + if n_gaussian < 1: + raise ValueError(f"`n_gaussian` must be positive, got {n_gaussian}") + if not np.isfinite(base_width) or base_width <= 0.0: + raise ValueError(f"`base_width` must be positive, got {base_width}") + if not np.isfinite(width_ratio) or width_ratio <= 1.0: + raise ValueError( + f"`width_ratio` must be greater than one, got {width_ratio}" + ) + width = base_width * np.power( + width_ratio, np.arange(n_gaussian, dtype=np.float64) + ) + if amplitude is None: + amplitude = np.sqrt(2.0 / np.pi) * np.log(width_ratio) / width + elif amplitude is None: + raise ValueError("`amplitude` is required with explicit `width`") + width = np.asarray(width, dtype=np.float64).reshape(-1) + if width.size < 1 or not np.all(np.isfinite(width)) or np.any(width <= 0.0): + raise ValueError("`width` must hold finite positive values") + amplitude = np.asarray(amplitude, dtype=np.float64).reshape(-1) + if amplitude.size == 1: + amplitude = np.full_like(width, amplitude[0]) + if amplitude.shape != width.shape or not np.all(np.isfinite(amplitude)): + raise ValueError( + "`amplitude` must be a finite scalar or hold one value per Gaussian" + ) + self.amplitude = amplitude.astype(self.prec) + self.log_width = np.log(width).astype(self.prec) + + def coefficients(self, like: Array) -> dict[str, Array]: + xp = array_api_compat.array_namespace(like) + return { + "amplitude": self._parameter("amplitude", like), + "inv_two_width_sq": 0.5 * xp.exp(-2.0 * self._parameter("log_width", like)), + } + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + xp = array_api_compat.array_namespace(r_sq) + gaussian = xp.exp(-r_sq[:, None] * coefficients["inv_two_width_sq"]) + return xp.sum(coefficients["amplitude"] * gaussian, axis=-1) + + def kspace(self, k_sq: Array) -> Array: + xp = array_api_compat.array_namespace(k_sq) + amplitude = self._parameter("amplitude", k_sq) + width = xp.exp(self._parameter("log_width", k_sq)) + # FT of a_k exp(-r^2 / 2 sigma_k^2) is a_k (2 pi)^{3/2} sigma_k^3 + # exp(-sigma_k^2 k^2 / 2); for the default amplitudes + # a_k = sqrt(2/pi) ln b / sigma_k the coefficient is 4 pi ln b sigma_k^2, + # whose ln-sigma trapezoid sum approaches 4 pi / k^2. The mode loop is + # unrolled: a (nk, M) intermediate breaks inductor's double backward of + # the k-space virial (NaN), while per-mode (nk,) terms compile fine. + terms = [] + for m in range(width.shape[0]): + weight = self._row_values(amplitude[m] * width[m] ** 3, k_sq) + rate = self._row_values(0.5 * width[m] ** 2, k_sq) + terms.append((2.0 * np.pi) ** 1.5 * weight * xp.exp(-rate * k_sq)) + return sum(terms[1:], terms[0]) + + def contact_value(self, like: Array) -> Array: + xp = array_api_compat.array_namespace(like) + amplitude = self._parameter("amplitude", like) + return xp.sum(amplitude) + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + # the slowest-decaying mode is the NARROWEST Gaussian: + # exp(-sigma_min^2 k^2 / 2) < tolerance + sigma_min = float( + np.exp(np.min(to_numpy_array(self.log_width).astype(np.float64))) + ) + return np.sqrt(2.0 * np.log(1.0 / tolerance)) / sigma_min + + @classmethod + def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: + return {"n_gaussian": int(np.asarray(variables["log_width"]).size)} diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py new file mode 100644 index 0000000000..3795fb8167 --- /dev/null +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -0,0 +1,730 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import math +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.common import ( + GLOBAL_NP_FLOAT_PRECISION, + NativeOP, + to_numpy_array, +) +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + K_GRID, + PME_MESH, + frame_id_from_n_node, + k_grid_bounds_of_cells, + segment_sum, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, +) +from .reciprocal import ( + kspace_energy, + mesh_energy, +) + +CHARGE_TARGET_SOURCES = ("neutral", "charge_spin", "fparam") +"""Sources of the total charge a constrained channel sums to.""" + +CHARGE_NET_NEURON = [64, 64] +"""Default hidden widths of the charge network.""" + +PME_ORDER = 8 +"""Default order of the B-splines that spread the charges to the mesh.""" + +PME_SPACING_FRACTION = 0.75 +"""Default mesh spacing, as a fraction of the Nyquist spacing of the wave-number cutoff. + +At order 8 the mesh evaluation of a periodic frame then stays within a few +parts in a million of the sum over the wave numbers, in the reciprocal energy +and in the force. +""" + +CHARGE_NET_SEED_INDEX = 10_000 +"""Child index of the fitting seed that seeds a charge network configured without a seed. + +The networks of the fitting and of the descriptor draw from small child +indices of their seeds, so the charge network draws a stream of its own even +when the fitting and descriptor seeds coincide. +""" + + +@BaseLongRange.register("latent_charge") +class LatentChargeLongRange(NativeOP, BaseLongRange): + r"""Energy of interacting latent charges, periodic or not. + + A charge network predicts :math:`C` latent charge channels + :math:`q_i^c` per atom, in units of :math:`e`, from the descriptor shared + with the energy fitting; padding and atoms owned by another rank carry no + charge. For a non-periodic batch the energy is the direct sum over the + pairs of atoms of each frame listed by the + :data:`~deepmd.dpmodel.utils.neighbor_graph.ALL_PAIRS` topology, each + unordered pair once; for a periodic batch (with ``periodic=True``) it is + the lattice sum of the kernel over the images of the cell, evaluated in + reciprocal space (see :mod:`~deepmd.dpmodel.long_range.reciprocal`), + + .. math:: + E_{\rm non-periodic} = k_e \sum_{i None: + if charge_net.vacuum_ref: + raise ValueError( + "the charge network reads the descriptor of the real atoms only " + "and cannot reference isolated atoms (`vacuum_ref`)" + ) + self.charge_net = charge_net + self.kernel = kernel + self.charge_constraint = self._normalize_constraint( + charge_constraint, charge_net + ) + self.periodic = bool(periodic) + self.kspace_tolerance = float(kspace_tolerance) + if ( + not np.isfinite(self.kspace_tolerance) + or not 0.0 < self.kspace_tolerance < 1.0 + ): + raise ValueError( + f"`kspace_tolerance` must be in (0, 1), got {self.kspace_tolerance}" + ) + # The frame charge/spin condition a compressed model is specialized + # to; see :meth:`fold_charge_state`. + self.charge_state = None + if k_grid_mode is None: + k_grid_mode = "pme" if self.periodic else "ragged" + self.k_grid_mode = str(k_grid_mode) + if self.k_grid_mode not in ("ragged", "fixed", "pme"): + raise ValueError( + f"`k_grid_mode` must be 'ragged', 'fixed' or 'pme', got {self.k_grid_mode!r}" + ) + if self.k_grid_mode in ("fixed", "pme") and not self.periodic: + raise ValueError( + f"`k_grid_mode='{self.k_grid_mode}'` requires `periodic=True`" + ) + if k_grid_bounds is not None and self.k_grid_mode != "fixed": + raise ValueError( + "`k_grid_bounds` is only meaningful with `k_grid_mode='fixed'`" + ) + self.decouple_k_grid = bool(decouple_k_grid) + if self.decouple_k_grid and not self.periodic: + raise ValueError("`decouple_k_grid` requires `periodic=True`") + # A decoupled grid reads the kernel once, at construction: the k-point + # selection stays at the initial kernel however training moves it, + # while the kernel values in the graph still use the live parameters. + self.frozen_k_max = ( + float(self.kernel.kspace_cutoff(self.kspace_tolerance)) + if self.decouple_k_grid + else None + ) + # The per-axis bounds of the shared integer grid of the fixed mode, + # zero until the data statistics size it (see + # :meth:`compute_input_stats`): an array of fixed shape, so that a + # backend keeping arrays in its state dict sees the same keys and + # shapes before and after the sizing, and a checkpoint restores the + # sized grid. + if k_grid_bounds is not None: + bounds = np.asarray(k_grid_bounds, dtype=np.int64).reshape(-1) + if bounds.shape != (3,) or np.any(bounds < 0) or not np.any(bounds): + raise ValueError( + "`k_grid_bounds` must be three non-negative integers, not " + f"all zero, got {list(k_grid_bounds)}" + ) + self.k_grid_bounds = bounds + elif self.k_grid_mode == "fixed": + self.k_grid_bounds = np.zeros(3, dtype=np.int64) + else: + self.k_grid_bounds = None + # The mesh of a periodic frame, on which training in the "pme" mode + # and every frozen model evaluate the reciprocal sum. + if not self.periodic and (pme_spacing is not None or pme_order is not None): + raise ValueError("`pme_spacing` and `pme_order` require `periodic=True`") + self.pme_order = PME_ORDER if pme_order is None else int(pme_order) + if self.pme_order < 2: + raise ValueError(f"`pme_order` must be at least 2, got {self.pme_order}") + if pme_spacing is not None: + pme_spacing = float(pme_spacing) + if not np.isfinite(pme_spacing) or pme_spacing <= 0.0: + raise ValueError( + f"`pme_spacing` must be finite and positive, got {pme_spacing}" + ) + self.pme_spacing = pme_spacing + + @staticmethod + def _normalize_constraint( + constraint: dict | None, charge_net: LatentChargeFitting + ) -> dict | None: + if constraint is None: + return None + source = str(constraint.get("source", "neutral")) + channel = int(constraint.get("channel", 0)) + fparam_index = int(constraint.get("fparam_index", 0)) + if source not in CHARGE_TARGET_SOURCES: + raise ValueError( + f"charge constraint source must be one of {CHARGE_TARGET_SOURCES}, " + f"got {source!r}" + ) + if not 0 <= channel < charge_net.dim_out: + raise ValueError( + f"constrained channel {channel} is out of range for " + f"{charge_net.dim_out} latent charge channel(s)" + ) + if source == "fparam" and not 0 <= fparam_index < charge_net.numb_fparam: + raise ValueError( + f"`fparam_index` {fparam_index} is out of range for " + f"{charge_net.numb_fparam} frame parameter(s)" + ) + return {"source": source, "channel": channel, "fparam_index": fparam_index} + + @classmethod + def from_config( + cls, + config: dict, + descriptor: Any, + fitting: Any, + fitting_seed: int | list[int] | None = None, + ) -> "LatentChargeLongRange": + """Build the module from its configuration. + + The charge network reads the descriptor and takes the conditioning + inputs of the energy fitting: its frame and atomic parameters, their + defaults and its case embedding. It follows the activation function, + skip connection and precision of the fitting unless its configuration + sets them, and its seed defaults to a child of the fitting seed. The + kernel parameters are stored in the precision of the charge network. + + Parameters + ---------- + config + The module configuration: ``channels``, ``charge_net``, ``kernel`` + (the ``les`` kernel when absent) and ``charge_constraint``. + descriptor + The descriptor of the model. + fitting + The energy fitting of the model. + fitting_seed + The seed of the energy fitting. + """ + charge_net_config = { + "neuron": CHARGE_NET_NEURON, + "activation_function": fitting.activation_function, + "resnet_dt": fitting.resnet_dt, + "precision": fitting.precision, + "seed": child_seed(fitting_seed, CHARGE_NET_SEED_INDEX), + **config.get("charge_net", {}), + } + kernel_config = { + "type": "les", + "precision": charge_net_config["precision"], + **config.get("kernel", {}), + } + kernel_cls = BaseInteractionKernel.get_class_by_type(kernel_config.pop("type")) + charge_net = LatentChargeFitting( + var_name="latent_charge", + ntypes=descriptor.get_ntypes(), + dim_descrpt=descriptor.get_dim_out(), + dim_out=int(config.get("channels", 1)), + numb_fparam=fitting.get_dim_fparam(), + numb_aparam=fitting.get_dim_aparam(), + dim_case_embd=fitting.dim_case_embd, + default_fparam=fitting.get_default_fparam(), + use_aparam_as_mask=fitting.use_aparam_as_mask, + mixed_types=descriptor.mixed_types(), + type_map=fitting.get_type_map(), + **charge_net_config, + ) + return cls( + charge_net, + kernel_cls(**kernel_config), + config.get("charge_constraint"), + periodic=config.get("periodic", False), + kspace_tolerance=config.get("kspace_tolerance", 1e-8), + k_grid_mode=config.get("k_grid_mode"), + k_grid_bounds=config.get("k_grid_bounds"), + decouple_k_grid=config.get("decouple_k_grid", False), + pme_spacing=config.get("pme_spacing"), + pme_order=config.get("pme_order"), + ) + + def _k_max(self) -> float: + """The wave-number cutoff selecting the k points of a periodic batch. + + A coupled grid reads the current kernel parameters, so a trained + kernel re-derives its grid every batch; a decoupled grid keeps the + cutoff of the initial parameters. + """ + if self.decouple_k_grid: + return self.frozen_k_max + return float(self.kernel.kspace_cutoff(self.kspace_tolerance)) + + def mesh_spacing(self) -> float: + r"""The largest spacing of the mesh of a periodic frame, in Å. + + The configured ``pme_spacing``, or :data:`PME_SPACING_FRACTION` of the + Nyquist spacing :math:`\pi / k_{max}` of the wave-number cutoff. + """ + if self.pme_spacing is not None: + return self.pme_spacing + return PME_SPACING_FRACTION * math.pi / self._k_max() + + def use_mesh_evaluation(self) -> None: + """Evaluate the reciprocal sum of every periodic batch on the mesh. + + A frozen model takes the mesh of every frame as an input of dynamic + size, sized by the cell of the frame, so that it serves any cell at + the accuracy of :meth:`mesh_spacing`; the explicit wave numbers of + the other modes have a count that no input layout fixes. The spacing + is fixed at the current kernel parameters, a constant of the frozen + model. + """ + if self.periodic: + self.pme_spacing = self.mesh_spacing() + self.k_grid_mode = "pme" + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the reciprocal-space topology of a periodic batch. + + The wave-number cutoff reads the kernel parameters at build time, so + that a trained kernel moves the k grid and the mesh with it, unless + ``decouple_k_grid`` keeps the initial cutoff. The k grid holds the + wave numbers of every frame within the cutoff, or the points of the + fixed integer grid with those beyond it masked; the mesh resolves + every cell of the batch at :meth:`mesh_spacing`. + + Raises + ------ + ValueError + If the fixed grid has not been sized. + """ + if not self.periodic: + return {} + if self.k_grid_mode == "pme": + return {PME_MESH: {"spacing": self.mesh_spacing()}} + params = {"k_max": self._k_max()} + if self.k_grid_mode == "fixed": + if not self._k_grid_sized(): + raise ValueError( + "the fixed k grid is not sized yet: run the data " + "statistics (`compute_input_stats`) on periodic " + "frames, or configure `k_grid_bounds`" + ) + params["bounds"] = to_numpy_array(self.k_grid_bounds) + return {K_GRID: params} + + def required_topologies(self) -> tuple[str, ...]: + if self.periodic: + if self.k_grid_mode == "pme": + return (ALL_PAIRS, PME_MESH) + return (ALL_PAIRS, K_GRID) + return (ALL_PAIRS,) + + def uses_charge_spin(self) -> bool: + # A folded state stands in for the condition. + return ( + self.charge_constraint is not None + and self.charge_constraint["source"] == "charge_spin" + and self.charge_state is None + ) + + def fold_charge_state(self, charge_spin: list[float]) -> None: + """Specialize the constraint to one frame charge/spin condition. + + A compressed descriptor folds the condition into its tables and reads + none at run time, so a constraint that reads its total charge from the + condition takes it from this state instead. The state is a constant + of the deployed model that a deployment rewrites, together with the + tables of the descriptor, when it serves another condition. + + Parameters + ---------- + charge_spin + The condition ``[charge, spin multiplicity]``. + """ + self.charge_state = np.asarray( + charge_spin, dtype=GLOBAL_NP_FLOAT_PRECISION + ).reshape(1, 2) + + def charge_state_constants(self) -> list[Array]: + return [] if self.charge_state is None else [self.charge_state] + + def rebuild_charge_state(self, charge_spin: Array) -> list[Array]: + if self.charge_state is None: + return [] + xp = array_api_compat.array_namespace(charge_spin) + state = xp.astype(charge_spin, self.charge_state.dtype) + return [xp.reshape(state, (1, 2))] + + def call(self, *args: Any, **kwargs: Any) -> Array: + return self.call_graph(*args, **kwargs) + + def call_graph( + self, + descriptor: Array, + atype: Array, + node_mask: Array, + n_node: Array, + topologies: dict[str, Any], + fparam: Array | None = None, + aparam: Array | None = None, + charge_spin: Array | None = None, + ) -> Array: + xp = array_api_compat.array_namespace(descriptor) + n = descriptor.shape[0] + pairs = topologies[ALL_PAIRS] + frame_id = frame_id_from_n_node(n_node, n_total=n) + fparam_node = None if fparam is None else xp.take(fparam, frame_id, axis=0) + charge = self.charge_net.call_graph( + descriptor, atype, fparam=fparam_node, aparam=aparam + )["latent_charge"] + # The pair sum cancels between charges of opposite sign, so it runs in + # the precision of the geometry. + charge = xp.astype(charge, pairs.vec.dtype) + charge = xp.where(node_mask[:, None], charge, xp.zeros_like(charge)) + if self.charge_constraint is not None: + charge = self._constrain( + charge, node_mask, frame_id, n_node, fparam, charge_spin + ) + energy = self._pair_branch(charge, pairs, n) + if self.periodic: + # The lattice sum of a periodic frame and its self term cancel to + # a small fraction of either, so the reciprocal-space energy is + # formed in float64 whatever the precision of the geometry. + charge_64 = xp.astype(charge, xp.float64) + if self.k_grid_mode == "pme": + reciprocal = mesh_energy( + charge_64, + topologies[PME_MESH], + frame_id, + self.kernel, + self.pme_order, + ) + else: + reciprocal = kspace_energy(charge_64, topologies[K_GRID], self.kernel) + energy = energy + xp.astype(reciprocal, energy.dtype) + return xp.reshape(energy, (n, 1)) + + def _pair_branch(self, charge: Array, pairs: Any, n: int) -> Array: + """Direct pair sum over the non-periodic frames, per atom.""" + xp = array_api_compat.array_namespace(charge) + src = pairs.index[0] + dst = pairs.index[1] + pair_charge = xp.sum( + xp.take(charge, src, axis=0) * xp.take(charge, dst, axis=0), axis=-1 + ) + r_sq = xp.sum(pairs.vec * pairs.vec, axis=-1) + pair_energy = pair_charge * self._pair_kernel(r_sq, src, n) + half_energy = (0.5 * COULOMB_CONSTANT) * xp.where( + pairs.mask, pair_energy, xp.zeros_like(r_sq) + ) + return segment_sum(half_energy, src, n) + segment_sum(half_energy, dst, n) + + def _pair_kernel(self, r_sq: Array, src: Array, n: int) -> Array: + """Kernel values of the pairs. + + The pairs are located on the node axis for a backend that brings the + kernel coefficients to the pairs through it. + + Parameters + ---------- + r_sq + Squared pair distances with shape ``(P,)``, in Ų. + src + First node of every pair with shape ``(P,)``. + n + The number of nodes. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + return self.kernel(r_sq) + + def _constrain( + self, + charge: Array, + node_mask: Array, + frame_id: Array, + n_node: Array, + fparam: Array | None, + charge_spin: Array | None, + ) -> Array: + """Shift the constrained channel so that it sums to the frame total charge.""" + xp = array_api_compat.array_namespace(charge) + nf = n_node.shape[0] + channel = self.charge_constraint["channel"] + target = self._total_charge(nf, charge, fparam, charge_spin) + weight = xp.astype(node_mask, charge.dtype) + count = segment_sum(weight, frame_id, nf) + total = segment_sum(charge[:, channel], frame_id, nf) + # A frame without a charged atom has nothing to shift. + shift = (total - target) / xp.maximum(count, xp.ones_like(count)) + constrained = charge[:, channel] - xp.take(shift, frame_id, axis=0) * weight + return xp.concat( + [charge[:, :channel], constrained[:, None], charge[:, channel + 1 :]], + axis=-1, + ) + + def _total_charge( + self, + nf: int, + like: Array, + fparam: Array | None, + charge_spin: Array | None, + ) -> Array: + """The total charge of every frame, which the constrained channel sums to. + + Parameters + ---------- + nf + The number of frames. + like + An array giving the namespace, dtype and device of the result. + fparam, charge_spin + As in :meth:`call_graph`. + + Returns + ------- + Array + The total charges with shape ``(nf,)``, in units of :math:`e`. + """ + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + source = self.charge_constraint["source"] + if source == "charge_spin" and self.charge_state is not None: + state = xp.asarray(self.charge_state, device=device) + target = xp.broadcast_to(state[:, 0], (nf,)) + elif source == "charge_spin": + if charge_spin is None: + raise ValueError( + "the charge constraint reads the total charge from " + "`charge_spin`, which was not given" + ) + target = charge_spin[:, 0] + elif source == "fparam": + if fparam is None: + raise ValueError( + "the charge constraint reads the total charge from " + "`fparam`, which was not given" + ) + target = fparam[:, self.charge_constraint["fparam_index"]] + else: + target = xp.zeros((nf,), dtype=like.dtype, device=device) + return xp.astype(target, like.dtype) + + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: + if self.k_grid_mode == "fixed" and not self._k_grid_sized(): + # Materialize the samples once for the grid sizing and the + # networks; a lazy sampler is shared by the callers. + sampled = merged() if callable(merged) else merged + self._size_k_grid(sampled) + merged = sampled + self.charge_net.compute_input_stats( + merged, protection=protection, stat_file_path=stat_file_path + ) + + def _size_k_grid(self, sampled: list[dict]) -> None: + r"""Size the fixed integer k grid by the largest sampled cell. + + Every sampled frame needs :math:`|n_d| \le k_{max} |a_d| / 2\pi` per + axis at the wave-number cutoff of the kernel; the grid takes the + elementwise maximum of the bounds over the frames. The kernel + parameters are read at statistics time, so the grid is sized for the + kernel of that moment; frames that outgrow it later, by larger cells + or a kernel whose cutoff grows, are truncated to the grid. + """ + k_max = self._k_max() + bounds = None + for frame in sampled: + box = frame.get("box") + if box is None: + continue + box = np.asarray(to_numpy_array(box), dtype=np.float64) + if box.size == 0: + continue + frame_max = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + bounds = frame_max if bounds is None else np.maximum(bounds, frame_max) + if bounds is None: + raise ValueError( + "the fixed k grid takes its size from the data statistics, " + "but the sampled frames hold no simulation cell; provide " + "periodic training data or configure `k_grid_bounds`" + ) + self.k_grid_bounds = bounds.astype(np.int64) + + def _k_grid_sized(self) -> bool: + """Whether the fixed integer grid has its bounds.""" + return bool(np.any(to_numpy_array(self.k_grid_bounds))) + + def change_type_map(self, type_map: list[str]) -> None: + self.charge_net.change_type_map(type_map=type_map) + + def set_case_embd(self, case_idx: int) -> None: + self.charge_net.set_case_embd(case_idx) + + def serialize(self) -> dict: + data = { + "@class": "LongRange", + "type": "latent_charge", + "@version": 1, + "charge_net": self.charge_net.serialize(), + "kernel": self.kernel.serialize(), + "charge_constraint": self.charge_constraint, + } + if self.periodic: + data["periodic"] = True + data["kspace_tolerance"] = self.kspace_tolerance + data["k_grid_mode"] = self.k_grid_mode + if self.k_grid_mode == "fixed" and self._k_grid_sized(): + data["k_grid_bounds"] = to_numpy_array(self.k_grid_bounds).tolist() + if self.decouple_k_grid: + data["decouple_k_grid"] = True + data["frozen_k_max"] = self.frozen_k_max + if self.periodic: + data["pme_order"] = self.pme_order + if self.pme_spacing is not None: + data["pme_spacing"] = self.pme_spacing + if self.charge_state is not None: + data["charge_state"] = to_numpy_array(self.charge_state).tolist() + return data + + @classmethod + def deserialize(cls, data: dict) -> "LatentChargeLongRange": + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + obj = cls( + LatentChargeFitting.deserialize(data["charge_net"]), + BaseInteractionKernel.deserialize(data["kernel"]), + data["charge_constraint"], + periodic=data.get("periodic", False), + kspace_tolerance=data.get("kspace_tolerance", 1e-8), + # A periodic record without a mode is one of the explicit wave + # numbers. + k_grid_mode=data.get("k_grid_mode", "ragged"), + k_grid_bounds=data.get("k_grid_bounds"), + decouple_k_grid=data.get("decouple_k_grid", False), + pme_spacing=data.get("pme_spacing"), + pme_order=data.get("pme_order"), + ) + if data.get("frozen_k_max") is not None: + # The constructor derives the frozen cutoff from the (possibly + # trained) kernel; the record holds the value the grid was + # actually built with. + obj.frozen_k_max = float(data["frozen_k_max"]) + if data.get("charge_state") is not None: + obj.fold_charge_state(np.asarray(data["charge_state"]).reshape(-1).tolist()) + return obj diff --git a/deepmd/dpmodel/long_range/reciprocal.py b/deepmd/dpmodel/long_range/reciprocal.py new file mode 100644 index 0000000000..09a574e46e --- /dev/null +++ b/deepmd/dpmodel/long_range/reciprocal.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Reciprocal-space energy of latent charges in periodic cells. + +For the charges :math:`q_i^c` of a periodic cell of volume :math:`V` and a +pair kernel :math:`K` with Fourier transform :math:`\hat K`, the energy of the +lattice of images, + +.. math:: + E = k_e \Big[ \frac{1}{2V} \sum_{\mathbf{k} \neq \mathbf{0}} \sum_c + \hat K(\mathbf{k}) \, |\hat\rho_c(\mathbf{k})|^2 + - \frac{1}{2} K(0) \sum_{i,c} (q_i^c)^2 \Big], + \qquad + \hat\rho_c(\mathbf{k}) = \sum_i q_i^c e^{i \mathbf{k} \cdot \mathbf{r}_i}, + +is the sum over the wave numbers of the cell minus the interaction of every +charge with itself. The :math:`\mathbf{k} = \mathbf{0}` mode is dropped, the +tin-foil convention, which for a channel with a net charge amounts to a +uniform neutralizing background. The energy is attributed to the atoms as +:math:`E_i = \frac{1}{2} \sum_c q_i^c \phi_i^c`, with the potential +:math:`\phi_i^c = \partial E / \partial q_i^c` that the other charges and the +images create at atom :math:`i`, as the pair sum of a non-periodic system +splits every pair term evenly between its atoms. + +Two evaluations serve the topologies of +:mod:`~deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph`: +:func:`kspace_energy` sums over explicit wave numbers, and +:func:`mesh_energy` evaluates the same sum by the smooth particle-mesh Ewald +method on a mesh. +""" + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, + xp_add_at, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, + MeshGrid, +) + +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, +) + + +def _self_potential( + kernel: BaseInteractionKernel, charge: Array, periodic_node: Array +) -> Array: + """The potential of every charge at itself, k_e K(0) q, on periodic nodes.""" + xp = array_api_compat.array_namespace(charge) + return ( + COULOMB_CONSTANT + * kernel.contact_value(charge) + * charge + * xp.astype(periodic_node, charge.dtype)[:, None] + ) + + +def _strained_rows(xp: Any, rows: Array, strain: Array) -> Array: + r"""Wave numbers of the strained cells, to first order in the strain. + + The strain :math:`\varepsilon` of a frame maps its cell :math:`B` to + :math:`B S^\mathsf{T}`, :math:`S = I + \varepsilon`, and a wave number, + a row :math:`\mathbf{k}`, to :math:`\mathbf{k} S^{-1}`. The strain of a + topology is zero, and the energy is differentiated once with respect to + it, for the virial; the value and the first derivative at zero strain + are those of :math:`S^{-1} = I - \varepsilon`, which maps the rows here + and keeps linear algebra out of the energy and its derivatives. + + Parameters + ---------- + rows + Wave numbers or reciprocal lattice vectors with shape ``(nf, ..., 3)``. + strain + The strain of every frame with shape ``(nf, 3, 3)``. + + Returns + ------- + Array + The strained wave numbers, of the shape of ``rows``. + """ + eps = xp.reshape(strain, (strain.shape[0],) + (1,) * (rows.ndim - 2) + (3, 3)) + return rows - xp.sum(rows[..., :, None] * eps, axis=-2) + + +def _strained_volume(volume: Array, strain: Array) -> Array: + r"""Volumes of the strained cells, :math:`V \det S = V (1 + \mathrm{tr}\,\varepsilon)` to first order.""" + return volume * (1.0 + strain[:, 0, 0] + strain[:, 1, 1] + strain[:, 2, 2]) + + +def kspace_energy( + charge: Array, + k_grid: KGrid, + kernel: BaseInteractionKernel, +) -> Array: + r"""Energy of the periodic frames over explicit wave numbers, per atom. + + The structure factors of the wave numbers of every frame come from the + phases of its nodes, and the potential at a node from the same phases: + :math:`\phi_i = \frac{k_e}{V} \sum_{\mathbf{k}} \hat K(\mathbf{k}) + [\cos(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Re}\,\hat\rho + + \sin(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Im}\,\hat\rho] - k_e K(0) q_i`. + Every frame pairs the wave numbers of its row of :attr:`KGrid.vec` with + the nodes of its row of :attr:`KGrid.slot`, in products batched over the + frames and the blocks of their wave numbers. The geometry is strained by + :attr:`KGrid.strain` (see :class:`KGrid`). + + Parameters + ---------- + charge + Node charges with shape ``(N, C)``, zero on the nodes that carry none. + k_grid + The wave numbers and the nodes of the periodic frames. + kernel + The pair kernel. + + Returns + ------- + Array + The energy of every node with shape ``(N,)``, in eV. A periodic frame + whose wave numbers all lie beyond the cutoff keeps its self term; a + non-periodic batch, whose k grid holds no nodes, has zero energy. + """ + xp = array_api_compat.array_namespace(charge) + device = array_api_compat.device(charge) + dtype = charge.dtype + n_node, n_channel = charge.shape + nf, n_slot = k_grid.slot.shape + # === Step 1. The nodes of every frame, (nf, ns, 3) and (nf, ns, C) === + real = k_grid.slot >= 0 + index = xp.reshape(xp.where(real, k_grid.slot, xp.zeros_like(k_grid.slot)), (-1,)) + coord = xp.take(xp.astype(k_grid.coord, dtype), index, axis=0) + coord = xp.reshape(coord, (nf, n_slot, 3)) + charge_slot = xp.reshape(xp.take(charge, index, axis=0), (nf, n_slot, n_channel)) + charge_slot = charge_slot * xp.astype(real, dtype)[:, :, None] + # === Step 2. Strained wave numbers and volumes: k' = S^-T k, V' = V det S === + # Every product over the wave numbers of a frame broadcasts the per-frame + # factor over the blocks, (nf, nb, kb, .), so that its sum runs over the + # points of a block. + strain = xp.astype(k_grid.strain, dtype) + k_vec = xp.astype(k_grid.vec, dtype) + strained_k = _strained_rows(xp, k_vec, strain) + volume = _strained_volume(xp.astype(k_grid.volume, dtype), strain) + # === Step 3. Structure factors, (nf, nb, kb, C) === + # The phase k' . r' = k . r does not change under the strain. + phase = xp.matmul(k_vec, xp.matrix_transpose(coord)[:, None, :, :]) + cos_phase = xp.cos(phase) + sin_phase = xp.sin(phase) + rho_real = xp.matmul(cos_phase, charge_slot[:, None, :, :]) + rho_imag = xp.matmul(sin_phase, charge_slot[:, None, :, :]) + # === Step 4. Potential at the nodes of every frame, (nf, ns, C) === + weight = COULOMB_CONSTANT * kernel.kspace(xp.sum(strained_k * strained_k, axis=-1)) + weight = weight / volume[:, None, None] + weight = xp.where(k_grid.mask, weight, xp.zeros_like(weight))[..., None] + block_potential = xp.matmul(xp.matrix_transpose(cos_phase), weight * rho_real) + block_potential = block_potential + xp.matmul( + xp.matrix_transpose(sin_phase), weight * rho_imag + ) + potential = xp.sum(block_potential, axis=1) + # Every node of a periodic frame carries its self term, whether or not a + # wave number of the frame lies within the cutoff; the slots of a + # non-periodic batch hold no nodes and hence no charge. + potential = ( + potential - COULOMB_CONSTANT * kernel.contact_value(charge) * charge_slot + ) + slot_energy = 0.5 * xp.sum(charge_slot * potential, axis=-1) + # === Step 5. Energy on the node axis === + # A padding slot carries no charge and adds a zero to the first node. + energy = xp.zeros((n_node,), dtype=dtype, device=device) + return xp_add_at(energy, index, xp.reshape(slot_energy, (-1,))) + + +def bspline_integer_weights(order: int) -> np.ndarray: + r"""The cardinal B-spline of order ``order`` at :math:`1, 2, \dots, p - 1`. + + These are the weights of the assignment factor + :math:`D(\theta) = \sum_{l=0}^{p-2} M_p(l + 1) e^{i \theta l}` of the + smooth particle-mesh Ewald method. + """ + return np.asarray( + [ + sum( + (-1.0) ** j * math.comb(order, j) * max(l + 1.0 - j, 0.0) ** (order - 1) + for j in range(order + 1) + ) + / math.factorial(order - 1) + for l in range(order - 1) + ] + ) + + +def _bspline(xp: Any, x: Array, order: int) -> Array: + r"""The cardinal B-spline :math:`M_p(x)` of order :math:`p`, elementwise. + + :math:`M_p(x) = \frac{1}{(p-1)!} \sum_{j=0}^{p} (-1)^j \binom{p}{j} + \max(x - j, 0)^{p-1}`, a piecewise polynomial supported on + :math:`[0, p]` with :math:`p - 2` continuous derivatives. + """ + value = None + for j in range(order + 1): + term = ((-1.0) ** j * math.comb(order, j)) * xp.maximum( + x - j, xp.zeros_like(x) + ) ** (order - 1) + value = term if value is None else value + term + return value / math.factorial(order - 1) + + +def _inverse3(xp: Any, m: Array) -> Array: + """Inverse of a batch of 3x3 matrices, its columns the cross products of the rows over the determinant.""" + a, b, c = m[..., 0, :], m[..., 1, :], m[..., 2, :] + bc = xp.linalg.cross(b, c) + ca = xp.linalg.cross(c, a) + ab = xp.linalg.cross(a, b) + det = xp.sum(a * bc, axis=-1) + return xp.stack([bc, ca, ab], axis=-1) / det[..., None, None] + + +def _stencil( + xp: Any, frac: Array, dims: tuple[Any, Any, Any], order: int +) -> tuple[Array, Array]: + r"""The mesh points every node spreads to and their weights. + + Along axis :math:`d` the scaled coordinate :math:`u = n_d s_d` of the + fractional coordinate :math:`s_d \in [0, 1)` reaches the ``order`` points + :math:`\lfloor u \rfloor - p + 1, \dots, \lfloor u \rfloor` with the + weights :math:`M_p(u - m)`; the stencil of a node is the outer product of + its three axes. + + Returns + ------- + tuple[Array, Array] + The flat mesh indices and the weights, both with shape + ``(N, order**3)``. + """ + device = array_api_compat.device(frac) + n_node = frac.shape[0] + shift = xp.arange(order, dtype=frac.dtype, device=device)[None, :] + offset = xp.arange(order, dtype=xp.int64, device=device)[None, :] + index = None + weight = None + for d, n_d in enumerate(dims): + u = frac[:, d] * n_d + base = xp.floor(u) + # The distance of point base - p + 1 + j from u is t + p - 1 - j. + w = _bspline(xp, (u - base)[:, None] + (order - 1) - shift, order) + m = (xp.astype(base, xp.int64)[:, None] - (order - 1) + offset) % n_d + if index is None: + index, weight = m, w + else: + index = xp.reshape(index[:, :, None] * n_d + m[:, None, :], (n_node, -1)) + weight = xp.reshape(weight[:, :, None] * w[:, None, :], (n_node, -1)) + return index, weight + + +def _assignment_factor( + m_axes: list[Array], dims: tuple[Any, Any, Any], order: int +) -> list[Array]: + r"""The inverse squared assignment factor :math:`|D_d(m)|^{-2}` of every axis. + + The B-spline spreading multiplies every mode of the charge mesh by + :math:`\prod_d D_d(m_d)`, :math:`D_d(m) = \sum_l M_p(l + 1) + e^{2\pi i m l / n_d}`, which the influence function divides out. + """ + xp = array_api_compat.array_namespace(m_axes[0]) + weights = bspline_integer_weights(order) + factors = [] + for m, n_d in zip(m_axes, dims, strict=True): + theta = (2.0 * math.pi) * m / n_d + real = xp.zeros_like(theta) + imag = xp.zeros_like(theta) + for l, w in enumerate(weights): + real = real + float(w) * xp.cos(theta * l) + imag = imag + float(w) * xp.sin(theta * l) + factors.append(1.0 / (real * real + imag * imag)) + return factors + + +def mesh_energy( + charge: Array, + mesh: MeshGrid, + frame_id: Array, + kernel: BaseInteractionKernel, + order: int, +) -> Array: + r"""Energy of the periodic frames by the smooth particle-mesh Ewald method, per atom. + + The charges spread to the mesh of every frame with B-splines of order + ``order`` along the fractional coordinates; a real FFT gives the half + spectrum :math:`Q(\mathbf{m})` of the charge mesh. The influence function + + .. math:: + G(\mathbf{m}) = \frac{k_e}{V} \hat K(\mathbf{k}_\mathbf{m}) + \prod_d |D_d(m_d)|^{-2}, + \qquad \mathbf{k}_\mathbf{m} = 2\pi B^{-1} \mathbf{m}, + + corrects for the spreading, so that the potential on the mesh is the + inverse transform of :math:`G Q` and the potential at a node its + interpolation by the same B-splines. :math:`|\mathbf{k}|^2` is the + quadratic form :math:`(2\pi)^2 \mathbf{m}^\mathsf{T} W \mathbf{m}` of the + reciprocal metric :math:`W = A^\mathsf{T} A`, :math:`A = S^{-\mathsf{T}} + B^{-1}` of the strained cell, and the strain scales the volume by + :math:`\det S`; the fractional coordinates do not depend on it. The zero + mode carries no weight, nor do the modes at the Nyquist frequency of an + axis, which have no unique wave number. + + Parameters + ---------- + charge + Node charges with shape ``(N, C)``, zero on the nodes that carry none. + mesh + The cells and the mesh of the periodic frames. + frame_id + The frame of every node with shape ``(N,)``. + kernel + The pair kernel. + order + The order of the B-splines. + + Returns + ------- + Array + The energy of every node with shape ``(N,)``, in eV; zero in a frame + masked out by :attr:`MeshGrid.mask`. + """ + xp = array_api_compat.array_namespace(charge) + device = array_api_compat.device(charge) + dtype = charge.dtype + nf = mesh.box.shape[0] + n_node, n_channel = charge.shape + m_full = [xp.astype(m, dtype) for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)] + n1, n2, n3 = (m.shape[0] for m in m_full) + n_mesh = n1 * n2 * n3 + box = xp.astype(mesh.box, dtype) + volume = xp.astype(mesh.volume, dtype) + # === Step 1. Stencils of the nodes on the mesh === + box_inv = _inverse3(xp, box) + # frac_d = sum_j r_j (B^-1)_jd, elementwise rather than a batched product + # of 3x3 matrices. + frac = xp.sum( + xp.astype(mesh.coord, dtype)[:, :, None] * xp.take(box_inv, frame_id, axis=0), + axis=1, + ) + frac = frac - xp.floor(frac) + index, weight = _stencil(xp, frac, (n1, n2, n3), order) + # Flat targets on the frame- and channel-major meshes, (N, C, order**3). + channel = xp.arange(n_channel, dtype=xp.int64, device=device) + first = (frame_id[:, None] * n_channel + channel[None, :]) * n_mesh + target = first[:, :, None] + index[:, None, :] + # === Step 2. Charge mesh and its half spectrum === + q_mesh = xp.zeros((nf * n_channel * n_mesh,), dtype=dtype, device=device) + q_mesh = xp_add_at( + q_mesh, + xp.reshape(target, (-1,)), + xp.reshape(charge[:, :, None] * weight[:, None, :], (-1,)), + ) + q_mesh = xp.reshape(q_mesh, (nf * n_channel, n1, n2, n3)) + spectrum = xp.fft.rfftn(q_mesh, axes=(-3, -2, -1)) + # === Step 3. Influence function on the half spectrum, (nf, n1, n2, h3) === + m3 = m_full[2][: n3 // 2 + 1] + m_half = [m_full[0], m_full[1], m3] + strain = xp.astype(mesh.strain, dtype) + # The rows of A^T = B^-T S^-1 are the reciprocal vectors of the strained + # cell, strained like wave numbers; W = A^T A. + a_t = _strained_rows(xp, xp.matrix_transpose(box_inv), strain) + metric = xp.sum(a_t[:, :, None, :] * a_t[:, None, :, :], axis=-1) + volume = _strained_volume(volume, strain) + g1 = m_half[0][:, None, None] + g2 = m_half[1][None, :, None] + g3 = m_half[2][None, None, :] + # Six monomials of the mode indices times six entries of the metric of + # every frame, elementwise. + monomials = (g1 * g1, g2 * g2, g3 * g3, g1 * g2, g1 * g3, g2 * g3) + entries = ( + metric[:, 0, 0], + metric[:, 1, 1], + metric[:, 2, 2], + 2.0 * metric[:, 0, 1], + 2.0 * metric[:, 0, 2], + 2.0 * metric[:, 1, 2], + ) + k_sq = None + for entry, monomial in zip(entries, monomials, strict=True): + term = entry[:, None, None, None] * monomial[None, ...] + k_sq = term if k_sq is None else k_sq + term + k_sq = (2.0 * math.pi) ** 2 * k_sq + # The zero mode and the Nyquist planes carry no weight. + keep = ( + (2.0 * xp.abs(g1) != n1) & (2.0 * xp.abs(g2) != n2) & (2.0 * xp.abs(g3) != n3) + ) + keep = keep & ((g1 != 0.0) | (g2 != 0.0) | (g3 != 0.0)) + keep = keep[None, ...] & mesh.mask[:, None, None, None] + safe_k_sq = xp.where(keep, k_sq, xp.ones_like(k_sq)) + shape = safe_k_sq.shape + # One row per frame and first mode index, along which the gradients of + # the kernel parameters sum first. + rows = xp.reshape(safe_k_sq, (nf * n1, n2 * (n3 // 2 + 1))) + transform = xp.reshape(kernel.kspace(rows), shape) + f1, f2, f3 = _assignment_factor(m_half, (n1, n2, n3), order) + factor = f1[:, None, None] * f2[None, :, None] * f3[None, None, :] + influence = (COULOMB_CONSTANT * transform * factor[None, ...]) / volume[ + :, None, None, None + ] + influence = xp.where(keep, influence, xp.zeros_like(influence)) + # === Step 4. Potential on the mesh and at the nodes === + spectrum = xp.reshape(spectrum, (nf, n_channel, n1, n2, n3 // 2 + 1)) + potential_mesh = xp.fft.irfftn( + spectrum * influence[:, None, ...], + s=(n1, n2, n3), + axes=(-3, -2, -1), + norm="forward", + ) + stencil_potential = xp.take( + xp.reshape(potential_mesh, (-1,)), xp.reshape(target, (-1,)), axis=0 + ) + potential = xp.sum( + xp.reshape(stencil_potential, (n_node, n_channel, -1)) * weight[:, None, :], + axis=-1, + ) + periodic_node = xp.take(mesh.mask, frame_id, axis=0) + potential = potential - _self_potential(kernel, charge, periodic_node) + return 0.5 * xp.sum(charge * potential, axis=-1) diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 7f3f8b9705..b5549d4adb 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -73,3 +73,5 @@ def compute_fitting_input_stats( The lazy sampled function to get data frames from different data systems. """ self.get_fitting_net().compute_input_stats(sampled_func) + if self.atomic_model.has_long_range(): + self.atomic_model.long_range.compute_input_stats(sampled_func) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b74fb6b697..4788e27b29 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -48,6 +48,7 @@ NeighborGraph, build_neighbor_graph, build_neighbor_graph_ase, + cast_topology_geometry, compact_nodes, expand_node_values, ) @@ -564,12 +565,20 @@ def _call_common_graph( # atom count the mask is all true and this is a renumbering by the # identity. ng, node_index = compact_nodes(ng, atype_flat >= 0) + atype_real = xp.take(atype_flat, node_index, axis=0) + topologies = self.atomic_model.build_topologies( + xp.take(xp.reshape(cc, (n_padded, 3)), node_index, axis=0), + ng.n_node, + bb, + atype=atype_real, + pair_excl=pair_excl, + ) # OUTPUT-AGNOSTIC standard model dict (````, ``_redu``, # derivative name-holders ``None``, plus int ``mask``), like the # dense ``call_common``. ``call_lower_graph`` masks virtual atoms # (atype<0) and sets the real int mask. model_predict = self.call_lower_graph( - atype=xp.take(atype_flat, node_index, axis=0), + atype=atype_real, n_node=ng.n_node, edge_index=ng.edge_index, edge_vec=ng.edge_vec, @@ -589,6 +598,7 @@ def _call_common_graph( else None ), charge_spin=charge_spin, + topologies=topologies, ) # Public ABI is rectangular (nf, nloc, *); the lower is flat over # the real atoms. Scatter per-atom keys back onto the padded width @@ -731,6 +741,7 @@ def forward_common_atomic_graph( comm_dict: dict | None = None, charge_spin: Array | None = None, spin: Array | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, Array]: """Model-level graph forward (no type cast). Analogue of the dense :meth:`forward_common_atomic`. @@ -773,6 +784,9 @@ def forward_common_atomic_graph( Per-node spin vectors, flat (N, 3), or ``None``. Forwarded unchanged to the atomic model's ``forward_common_atomic_graph`` (and, from there, the descriptor's ``call_graph``). + topologies + The topologies the atomic model declares by + ``required_topologies()``, on the node axis of the graph. Returns ------- @@ -795,6 +809,7 @@ def forward_common_atomic_graph( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) return fit_output_to_model_output_graph( atomic_ret, @@ -817,6 +832,7 @@ def call_common_lower_graph( comm_dict: dict | None = None, charge_spin: Array | None = None, spin: Array | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, Array]: """Graph-native PUBLIC lower (dpa1/se_atten concat-tebd, attention included). @@ -858,6 +874,9 @@ def call_common_lower_graph( Per-node spin vectors, flat (N, 3), or ``None``. Cast to the model precision alongside the other node inputs and forwarded unchanged to :meth:`forward_common_atomic_graph`. + topologies + The topologies the atomic model declares by + ``required_topologies()``, on the node axis of the graph. Returns ------- @@ -871,6 +890,12 @@ def call_common_lower_graph( charge_spin=charge_spin, spin=spin, ) + if topologies: + # The topologies carry geometry like the edge vectors. + topologies = { + name: cast_topology_geometry(topology, edge_vec.dtype) + for name, topology in topologies.items() + } model_predict = self.forward_common_atomic_graph( atype, n_node, @@ -883,6 +908,7 @@ def call_common_lower_graph( comm_dict=comm_dict, charge_spin=cs, spin=sp, + topologies=topologies, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -926,6 +952,7 @@ def change_out_bias( self, merged: Any, bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -936,8 +963,17 @@ def change_out_bias( bias_adjust_mode : str The mode for changing output bias: 'change-by-statistic' or 'set-by-statistic'. + preset_out_bias : dict, optional + Per-call preset override. Omit to use the configured preset; + pass an empty dict to fit without fixing preset outputs. + Non-finite stored cells covered by the checkpoint preset are + recovered before residual fitting; finite values are retained. """ - self.atomic_model.change_out_bias(merged, bias_adjust_mode=bias_adjust_mode) + self.atomic_model.change_out_bias( + merged, + bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, + ) def _input_type_cast( self, diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 66d32f8f8b..21c3be799b 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -13,6 +13,9 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) +from deepmd.dpmodel.long_range import ( + BaseLongRange, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -47,6 +50,7 @@ atomic_model=DPAtomicModel, pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, + long_range_base=BaseLongRange, ) get_zbl_model = _model_factory.get_zbl_model diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 38a4f4b9e4..118b09a75f 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -65,6 +65,83 @@ def get_model_components( return descriptor, fitting, fitting_type +def get_long_range( + data: dict, + descriptor: Any, + fitting: Any, + *, + long_range_base: type | None, + backend_name: str, +) -> Any | None: + """Construct the long-range module a model configuration adds to its fitting. + + The module reads the descriptor of the model, and its networks take the + conditioning inputs of the energy fitting they accompany. + + Parameters + ---------- + data : dict + The model configuration, with the optional ``long_range`` section. + descriptor : Any + The descriptor of the model. + fitting : Any + The energy fitting of the model. + long_range_base : type or None + Backend long-range registry base class; ``None`` for a backend + without long-range modules. + backend_name : str + Backend name used in error messages. + + Returns + ------- + Any or None + The backend-native long-range module, or ``None`` when none is + configured. + + Raises + ------ + NotImplementedError + If a long-range module is configured for a backend without them. + """ + config = data.get("long_range") + if config is None: + return None + if long_range_base is None: + raise NotImplementedError( + f"long-range modules are not supported in the {backend_name} backend" + ) + config = copy.deepcopy(config) + long_range_cls = long_range_base.get_class_by_type(config.pop("type")) + return long_range_cls.from_config( + config, + descriptor, + fitting, + fitting_seed=data.get("fitting_net", {}).get("seed"), + ) + + +def _long_range_kwargs( + data: dict, + descriptor: Any, + fitting: Any, + *, + long_range_base: type | None, + backend_name: str, +) -> dict[str, Any]: + """The ``long_range`` constructor argument, absent when none is configured. + + Model classes without a long-range module keep their constructor contract. + """ + long_range = get_long_range( + data, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ) + return {} if long_range is None else {"long_range": long_range} + + def get_standard_model( data: dict, *, @@ -72,6 +149,7 @@ def get_standard_model( fitting_base: type, model_base: type, backend_name: str, + long_range_base: type | None = None, ) -> Any: """Construct a standard model through backend registries.""" descriptor, fitting, fitting_type = get_model_components( @@ -89,6 +167,13 @@ def get_standard_model( atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=data.get("pair_exclude_types", []), preset_out_bias=data.get("preset_out_bias"), + **_long_range_kwargs( + data, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ), ) @@ -103,6 +188,11 @@ def get_zbl_model( backend_name: str, ) -> Any: """Construct a ZBL model from backend-native atomic model classes.""" + if data.get("long_range") is not None: + raise NotImplementedError( + "a long-range module is not supported with the tabulated ZBL " + "potential (`use_srtab`)" + ) descriptor, fitting, fitting_type = get_model_components( data, descriptor_base=descriptor_base, @@ -141,6 +231,7 @@ def get_linear_atomic_model( pairtab_model: type, linear_atomic_model: type | None = None, descriptor_child_builder: "Callable[[dict], Any | None] | None" = None, + long_range_base: type | None = None, ) -> Any: """Build the ``LinearEnergyAtomicModel`` composition from a config. @@ -181,6 +272,9 @@ def get_linear_atomic_model( back to the generic registry build. Backends use it to route family-specific model types (e.g. DPA4/SeZM) through their validated builders. + long_range_base : type, optional + Backend long-range registry base class, which builds the long-range + module of a child; ``None`` for a backend without long-range modules. Raises ------ @@ -188,8 +282,11 @@ def get_linear_atomic_model( If more than one ``inner_potential`` child is given, if an ``inner_potential`` child has no unique learned sibling, if a bridged composition does not use ``weights: "sum"``, if a child - carries a ``bridging_method`` flag, or if a child is of an + carries a ``bridging_method`` flag, if a child with a long-range + module carries its own ``pair_exclude_types``, or if a child is of an unsupported kind. + NotImplementedError + If a child with a long-range module carries ``atom_exclude_types``. """ from deepmd.dpmodel.atomic_model.inner_potential import ( InnerPotentialAtomicModel, @@ -311,6 +408,24 @@ def get_linear_atomic_model( if descriptor_child_builder is not None: child = descriptor_child_builder(sub) if child is None: + # A long-range child is checked as the same model alone: the + # pair exclusion of the composition shapes the graph and the + # pair topology the children share, and atom exclusion is + # unsupported with a long-range module. + if sub.get("long_range") is not None: + if sub.get("pair_exclude_types"): + raise ValueError( + "a linear_ener sub-model with a long-range module " + "takes `pair_exclude_types` from the composition, " + "whose sub-models share one neighbor graph and one " + "pair topology; set it on the composition" + ) + if sub.get("atom_exclude_types"): + raise NotImplementedError( + "`atom_exclude_types` is not supported with a " + "long-range module, which attributes each pair " + "energy to both atoms of the pair" + ) descriptor, fitting, _ = get_model_components( sub, descriptor_base=descriptor_base, @@ -325,6 +440,13 @@ def get_linear_atomic_model( fitting, type_map=sub["type_map"], preset_out_bias=preset_out_bias, + **_long_range_kwargs( + sub, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ), ) built[i] = child else: @@ -467,6 +589,7 @@ def __init__( pairtab_model: type | None = None, zbl_model: type | None = None, linear_atomic_model: type | None = None, + long_range_base: type | None = None, ) -> None: """Store backend-native classes used by all model construction paths.""" self.descriptor_base = descriptor_base @@ -477,6 +600,7 @@ def __init__( self.pairtab_model = pairtab_model self.zbl_model = zbl_model self.linear_atomic_model = linear_atomic_model + self.long_range_base = long_range_base def get_model_components(self, data: dict) -> tuple[Any, Any, str]: """Construct descriptor and fitting objects for this backend.""" @@ -495,6 +619,17 @@ def get_standard_model(self, data: dict) -> Any: fitting_base=self.fitting_base, model_base=self.model_base, backend_name=self.backend_name, + long_range_base=self.long_range_base, + ) + + def get_long_range(self, data: dict, descriptor: Any, fitting: Any) -> Any | None: + """Construct the long-range module of a model config for this backend.""" + return get_long_range( + data, + descriptor, + fitting, + long_range_base=self.long_range_base, + backend_name=self.backend_name, ) def get_linear_atomic_model( @@ -515,6 +650,7 @@ def get_linear_atomic_model( pairtab_model=self.pairtab_model, linear_atomic_model=self.linear_atomic_model, descriptor_child_builder=descriptor_child_builder, + long_range_base=self.long_range_base, ) def get_zbl_model(self, data: dict) -> Any: diff --git a/deepmd/dpmodel/model/native_spin_model.py b/deepmd/dpmodel/model/native_spin_model.py index 5df83d91df..6fdd5744d6 100644 --- a/deepmd/dpmodel/model/native_spin_model.py +++ b/deepmd/dpmodel/model/native_spin_model.py @@ -79,6 +79,10 @@ class NSM(T_Model, NativeSpinModelKind): def __init__(self, *args: Any, spin: Spin, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + if self.atomic_model.has_long_range(): + raise NotImplementedError( + "native spin is not supported for a model with a long-range module" + ) self.spin = spin self.ntypes_real = self.spin.ntypes_real # Per-real-type 0/1 spin gate, derived from ``use_spin`` and hence diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 8678e51398..c1bb8c3fe1 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -31,6 +31,10 @@ from deepmd.dpmodel.output_def import ( ModelOutputDef, ) +from deepmd.utils.preset_out_bias import ( + normalize_preset_out_bias, + remap_preset_out_bias, +) from deepmd.utils.spin import ( Spin, ) @@ -118,6 +122,11 @@ def __init__( spin: Spin, ) -> None: super().__init__() + if backbone_model.atomic_model.has_long_range(): + raise NotImplementedError( + "the virtual-atom spin scheme is not supported for a model with " + "a long-range module" + ) self.backbone_model = backbone_model self.spin = spin # Spin graph-lower unsupported: carry-all routing diverges on @@ -603,6 +612,7 @@ def change_out_bias( self, merged: Callable[[], list[dict]] | list[dict], bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. @@ -620,13 +630,24 @@ def change_out_bias( 'change-by-statistic' : perform predictions on labels of target dataset, and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. + preset_out_bias : dict, optional + Preset on the public real-atom type map. It is normalized and + extended with unassigned virtual types before forwarding. Omit + to retain the model preset; an empty dict fits without fixing + preset outputs. Non-finite stored cells covered by a checkpoint + preset are recovered before residual fitting. """ + preset_out_bias = remap_preset_out_bias( + normalize_preset_out_bias(preset_out_bias, self.get_type_map()), + list(range(self.get_ntypes())) + [-1] * self.get_ntypes(), + ) spin_sampled_func = self._get_spin_sampled_func( merged if callable(merged) else lambda: merged ) self.backbone_model.change_out_bias( spin_sampled_func, bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, ) def change_type_map( diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index 1463dad928..25969593d9 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -12,6 +12,12 @@ make_task_maps, print_data_summaries, ) +from .data_progress import ( + DATA_PROGRESS_CHECKPOINT_KEY, + collect_training_data_progress, + preserve_training_data_progress, + restore_training_data_progress, +) from .entrypoint import ( AbstractTrainEntrypoint, TrainEntrypointOptions, @@ -44,6 +50,7 @@ ) __all__ = [ + "DATA_PROGRESS_CHECKPOINT_KEY", "DEFAULT_TASK_KEY", "AbstractTrainEntrypoint", "AbstractTrainer", @@ -64,9 +71,12 @@ "build_checkpoint_stores", "change_model_out_bias", "change_model_out_bias_by_task", + "collect_training_data_progress", "iter_training_task_configs", "make_task_maps", + "preserve_training_data_progress", "print_data_summaries", "resolve_keep_ckpt_count", "resolve_step_schedule", + "restore_training_data_progress", ] diff --git a/deepmd/dpmodel/train/data.py b/deepmd/dpmodel/train/data.py index 3e91dcfc53..86fde438f2 100644 --- a/deepmd/dpmodel/train/data.py +++ b/deepmd/dpmodel/train/data.py @@ -14,6 +14,14 @@ Any, ) +from deepmd.utils.data import ( + DataRequirementItem, + min_pair_dist_requirement, +) +from deepmd.utils.data_config import ( + DataConfig, + data_group_configs, +) from deepmd.utils.stat_file import ( StatFileMode, StatFileSpec, @@ -37,7 +45,7 @@ class TrainingTaskConfig: key: str model_params: Mapping[str, Any] - training_data_params: Mapping[str, Any] + training_data_params: DataConfig validation_data_params: Mapping[str, Any] | None stat_file: str | None valid_numb_batch: int @@ -110,6 +118,33 @@ def make_task_maps( return training_data, validation_data, stat_data +def add_training_data_requirements( + data: Any, + config: DataConfig, + model_params: Mapping[str, Any], + requirements: list[DataRequirementItem], + register: Callable[[Any, list[DataRequirementItem]], None], +) -> bool: + """Register each group's labels and its own derived pair-clearance margin. + + The returned flag enables margin filtering in the trainer and statistics + collector. Distances belong to individual groups' requirements, not to + that task-wide enable flag. + """ + groups = data_group_configs(config) + datasets = data.children if len(groups) > 1 else (data,) + filter_enabled = [] + for dataset, group in zip(datasets, groups, strict=True): + pair_filter = min_pair_dist_requirement( + model_params, float(group.get("min_pair_dist", 0.0)) + ) + register(dataset, requirements + ([] if pair_filter is None else [pair_filter])) + filter_enabled.append(pair_filter is not None) + if len(groups) > 1: + data.set_filter_enabled(filter_enabled) + return any(filter_enabled) + + def print_data_summaries( training_data: Mapping[str, Any], validation_data: Mapping[str, Any | None], diff --git a/deepmd/dpmodel/train/data_progress.py b/deepmd/dpmodel/train/data_progress.py new file mode 100644 index 0000000000..0c5858fe78 --- /dev/null +++ b/deepmd/dpmodel/train/data_progress.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Small, backend-independent checkpoints of the next training batches. + +Data sources own their cursors and deterministic sampler state. This module +adds task names and the shared data-source RNG, not decoded data or permutations. +""" + +from collections.abc import ( + Iterator, + Mapping, +) +from contextlib import ( + contextmanager, +) +from typing import ( + Any, +) + +from deepmd.utils import random as dp_random + +DATA_PROGRESS_CHECKPOINT_KEY = "data_progress" +DATA_PROGRESS_VERSION = 1 + + +def collect_training_data_progress( + training_data_by_task: Mapping[str, Any], +) -> dict[str, Any]: + """Collect supported data-source cursors and the shared data-source RNG.""" + return { + "version": DATA_PROGRESS_VERSION, + "tasks": { + key: data.state_dict() + for key, data in training_data_by_task.items() + if callable(getattr(data, "state_dict", None)) + and callable(getattr(data, "load_state_dict", None)) + }, + "rng": dp_random.get_state(), + } + + +def restore_training_data_progress( + training_data_by_task: Mapping[str, Any], + progress: Mapping[str, Any] | None, +) -> None: + """Restore cursors lazily; missing progress preserves legacy restarts. + + The source implementations validate their sampling configuration and must + not read or decode data here. Restoring the shared RNG last also isolates + it from any source-specific reconstruction. + """ + if progress is None: + return + if progress.get("version") != DATA_PROGRESS_VERSION: + raise ValueError("Unsupported data-progress checkpoint version.") + tasks = progress.get("tasks") + if not isinstance(tasks, Mapping): + raise ValueError("data-progress 'tasks' must be a mapping.") + supported = { + key: data + for key, data in training_data_by_task.items() + if callable(getattr(data, "load_state_dict", None)) + and callable(getattr(data, "state_dict", None)) + } + if set(tasks) != set(supported): + raise ValueError("data-progress tasks do not match the training data sources.") + if "rng" not in progress: + raise ValueError("data-progress checkpoint is missing the selection RNG.") + for key, data in supported.items(): + data.load_state_dict(tasks[key]) + dp_random.set_state(progress["rng"]) + + +@contextmanager +def preserve_training_data_progress( + training_data_by_task: Mapping[str, Any], +) -> Iterator[None]: + """Keep diagnostics and compilation from consuming the training stream.""" + progress = collect_training_data_progress(training_data_by_task) + try: + yield + finally: + restore_training_data_progress(training_data_by_task, progress) diff --git a/deepmd/dpmodel/train/schedule.py b/deepmd/dpmodel/train/schedule.py index 8fe2749dfc..1279343fc3 100644 --- a/deepmd/dpmodel/train/schedule.py +++ b/deepmd/dpmodel/train/schedule.py @@ -88,8 +88,9 @@ def resolve_step_schedule( Task keys in the order used by the trainer. Single-task runs pass the one key under which their data is registered. training_data : Mapping[str, Any] - Training data of every task, keyed like ``model_keys``. Only consulted - when multi-task sampling weights fall back to the per-task data size. + Training data of every task, keyed like ``model_keys``. Only consulted + when multi-task sampling weights default to native data sizes: + NumPy/HDF5 system counts or retained LMDB frame counts. epoch_length : Callable[[str], int] Maps a task key to the number of steps *this rank* performs during one epoch of that task. Backends whose data pipeline is sharded across diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 22cb543ced..20757fae9f 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -74,12 +74,17 @@ def change_model_out_bias( *, bias_adjust_mode: str = "change-by-statistic", recompute_input_stats: bool = False, + preset_out_bias: dict | None = None, ) -> Any: """Change one model's output bias and log the before/after values.""" old_bias = deepcopy(model.get_out_bias()) + preset_kwargs = ( + {"preset_out_bias": preset_out_bias} if preset_out_bias is not None else {} + ) model.change_out_bias( sample_func, bias_adjust_mode=bias_adjust_mode, + **preset_kwargs, ) new_bias = deepcopy(model.get_out_bias()) diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 0d7ea42a89..f5e39a98ac 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -5,7 +5,9 @@ Backend-specific wrappers (PyTorch Dataset, JAX, etc.) import from here. """ +import copy import dataclasses +import hashlib import logging import math import multiprocessing @@ -30,6 +32,9 @@ from dataclasses import ( dataclass, ) +from itertools import ( + islice, +) from pathlib import ( Path, ) @@ -60,6 +65,7 @@ DataRequirementItem, DataRequirementSourcePolicy, ) +from deepmd.utils.data_path import is_lmdb as is_lmdb log = logging.getLogger(__name__) @@ -657,16 +663,27 @@ def _scan_lmdb_path_in_worker( executor.shutdown(wait=True, cancel_futures=True) -def _remap_atom_types(atype: np.ndarray, type_remap: np.ndarray) -> np.ndarray: - """Remap real atom types while preserving negative virtual sentinels. - - Positive indices retain NumPy's normal bounds checking, so malformed LMDB - data cannot be silently reinterpreted as a different species. - """ - remapped_atype = atype.astype(np.int64, copy=True) - real_atom_mask = remapped_atype >= 0 - remapped_atype[real_atom_mask] = type_remap[remapped_atype[real_atom_mask]] - return remapped_atype +def _make_type_remap( + lmdb_type_map: list[str] | None, type_map: list[str] | None +) -> np.ndarray | None: + """Map declared source types to model indices, marking unsupported entries.""" + if ( + lmdb_type_map is None + or type_map is None + or list(lmdb_type_map) == list(type_map) + ): + return None + remap = np.array( + [type_map.index(name) if name in type_map else -1 for name in lmdb_type_map], + dtype=np.int32, + ) + log.info( + "Type remapping: LMDB %s -> model %s, remap=%s", + lmdb_type_map, + type_map, + remap, + ) + return remap @dataclass @@ -685,7 +702,8 @@ class LmdbDecodeConfig: natoms Fallback atom count for records without ``atom_types``. type_remap - Optional LMDB-type to model-type lookup table. + Optional LMDB-type to model-type lookup table. Negative entries denote + unsupported elements, rejected only when referenced by a real atom. data_requirements Registered data requirements keyed by field name. dataset @@ -699,6 +717,9 @@ class LmdbDecodeConfig: their inputs and derive their labels here, before the model runs, which is the only place that works for a backend whose loss never sees the model. ``None``, the default, leaves decoding unchanged. + lmdb_type_map + Source element symbols used to identify unsupported atom types in + frame-level diagnostics. """ ntypes: int @@ -708,6 +729,7 @@ class LmdbDecodeConfig: dataset: str = "" type_map: list[str] | None = None frame_transform: Callable[[dict[str, Any], int], dict[str, Any]] | None = None + lmdb_type_map: list[str] | None = None def _requirement_dtype(requirement: Any) -> np.dtype: @@ -1041,7 +1063,26 @@ def decode_lmdb_frame( if "atype" in frame and isinstance(frame["atype"], np.ndarray): frame["atype"] = frame["atype"].reshape(-1).astype(np.int64, copy=False) if config.type_remap is not None: - frame["atype"] = _remap_atom_types(frame["atype"], config.type_remap) + atype = frame["atype"] + real_atom_mask = atype >= 0 + real_types = atype[real_atom_mask] + # Negative input sentinels are not lookup indices. Positive indices + # retain NumPy's bounds checks for malformed source type indices. + mapped = config.type_remap[real_types] + unsupported = mapped < 0 + if np.any(unsupported): + indices = np.unique(real_types[unsupported]) + missing = ( + [config.lmdb_type_map[index] for index in indices] + if config.lmdb_type_map is not None + else indices.tolist() + ) + raise ValueError( + f"LMDB frame {original_key} in '{config.dataset}' uses atom " + f"types {missing} not found in model type_map {config.type_map}" + ) + frame["atype"] = atype.copy() + frame["atype"][real_atom_mask] = mapped if "virial" in frame and isinstance(frame["virial"], np.ndarray): frame["virial"] = ( frame["virial"] @@ -1726,7 +1767,14 @@ def __init__( self._reader = reader self._sampler = sampler self._epoch = 0 + # These name the next batch the consumer will see, not the pass that + # the one-ahead decoder (or the sampler itself) has already entered. + self._next_epoch = 0 + self._next_batch_index = 0 + self._resume_batch_index = 0 + self._restore_config: dict[str, Any] | None = None self._queued_pass_length = 0 + self._queued_nframes = 0.0 self._last_pass_length: int | None = None self._iterator = self._iter_epoch() self._num_workers = num_workers @@ -1747,13 +1795,19 @@ def __next__(self) -> dict[str, Any]: batch = self._collect(pending) elif self._deferred_indices is not None: indices, self._deferred_indices = self._deferred_indices, None - batch = self._reader.decode_batch(indices) + batch = self._decode(indices) else: - indices, self._queued_pass_length = self._next_indices() + indices, self._queued_pass_length, self._queued_nframes = ( + self._next_indices() + ) batch = self._decode(indices) self._last_pass_length = self._queued_pass_length - indices, self._queued_pass_length = self._next_indices() + self._next_batch_index += 1 + if self._next_batch_index == self._last_pass_length: + self._next_epoch += 1 + self._next_batch_index = 0 + indices, self._queued_pass_length, self._queued_nframes = self._next_indices() self._schedule(indices) return batch @@ -1833,7 +1887,7 @@ def _lose_the_pool(self, futures: "Sequence[Future[dict[str, Any]]]" = ()) -> No "decoders outlive the session that started it." ) - def _next_indices(self) -> tuple[list[int], int]: + def _next_indices(self) -> tuple[list[int], int, float]: """Return the next sampler batch and restart after exhaustion.""" try: return next(self._iterator) @@ -1842,8 +1896,11 @@ def _next_indices(self) -> tuple[list[int], int]: self._iterator = self._iter_epoch() return next(self._iterator) - def _iter_epoch(self) -> Iterator[tuple[list[int], int]]: - """Tag batches with the length of the sampler pass that owns them.""" + def _iter_epoch(self) -> Iterator[tuple[list[int], int, float]]: + """Snapshot each batch's pass length and rank-averaged frame count.""" + if self._restore_config is not None: + self._check_restore_config(self._restore_config) + self._restore_config = None set_epoch = getattr(self._sampler, "set_epoch", None) if callable(set_epoch): set_epoch(self._epoch) @@ -1851,8 +1908,148 @@ def _iter_epoch(self) -> Iterator[tuple[list[int], int]]: # enter another pass before the current batch reaches its consumer. # Snapshot the length lazily, after data requirements and layout are set. pass_length = len(self._sampler) - for indices in self._sampler: - yield indices, pass_length + if pass_length == 0 and self._resume_batch_index == 0: + return + if not 0 <= self._resume_batch_index < pass_length: + raise ValueError( + f"LMDB batch index {self._resume_batch_index} is out of range " + f"for epoch {self._epoch} with {pass_length} batches" + ) + # This iterator already owns epoch advancement. Retain the sampler's + # current pass instead of its standalone __iter__ advancing and + # invalidating it, so diagnostic reads can rewind within the pass. + pass_batches = ( + self._sampler.batches() + if isinstance(self._sampler, LmdbBatchSampler) + else self._sampler + ) + frame_counts = getattr(self._sampler, "batch_frame_counts", None) + counts = frame_counts() if frame_counts is not None else None + start = self._resume_batch_index + batches = islice(pass_batches, start, None) + self._resume_batch_index = 0 + for index, indices in enumerate(batches, start): + nframes = len(indices) if counts is None else counts[index] + yield indices, pass_length, float(nframes) + + @property + def uniform_sampling_mass(self) -> float: + """Weight the next batch by its retained systems per sampled frame. + + Only sampler indices are inspected; no extra batch is decoded. The + frame count belongs to the queued batch, even when prefetching has + already advanced the sampler to another pass. Distributed samplers + report the same rank-averaged count on every rank. + """ + nsystems = self._reader.sampling_nsystems + if not nsystems: + return 0.0 + if self._pending is None and self._deferred_indices is None: + ( + self._deferred_indices, + self._queued_pass_length, + self._queued_nframes, + ) = self._next_indices() + return nsystems / self._queued_nframes + + @property + def uniform_nominal_mass(self) -> float: + """Return the systems-per-frame weight averaged over a sampler pass.""" + nsystems = self._reader.sampling_nsystems + if not nsystems: + return 0.0 + return nsystems / float(np.mean(self._sampler.batch_frame_counts())) + + def _sampling_config(self) -> dict[str, Any]: + """Describe sampling without building a pass or reading any frames.""" + reader, sampler = self._reader, self._sampler + if not isinstance(sampler, (LmdbBatchSampler, DistributedLmdbBatchSampler)): + raise TypeError("LMDB progress requires an LMDB batch sampler") + return { + "path": reader.lmdb_path, + "transaction": reader._transaction().id(), + "nframes": len(reader), + "type_map": list(reader.type_map), + "batch_rule": [ + reader.batch_size, + reader._auto_rule, + reader._max_rule, + reader._filter_rule, + reader._mix_rule, + ], + "ragged": reader.ragged_batches, + "availability_keys": list( + _availability_requirement_keys(reader._data_requirements) + ), + "shuffle": sampler._shuffle, + "configured_seed": sampler._configured_seed, + "world_size": getattr(sampler, "world_size", 1), + "block_targets": hashlib.sha256( + msgpack.packb(sampler._block_targets, use_bin_type=True) + ).hexdigest(), + } + + def _check_restore_config( + self, config: dict[str, Any], *, final: bool = True + ) -> None: + """Reject changed sampling, deferring trainer-selected configuration.""" + current = self._sampling_config() + for key in current: + if not final and key in {"ragged", "availability_keys"}: + continue + if config.get(key) != current[key]: + raise ValueError(f"LMDB data progress is incompatible: {key} changed") + + def state_dict(self) -> dict[str, Any]: + """Return a compact next-batch checkpoint for this rank. + + No decoded frames, permutations or prefetched work are serialized. + Each rank restores its deterministic partition using its own consumed + cursor. Filtering can advance ranks differently, so distributed + checkpoint writers must retain every rank's state. + """ + config = ( + self._restore_config + if self._restore_config is not None + else self._sampling_config() + ) + return { + "version": 1, + "epoch": self._next_epoch, + "batch_index": self._next_batch_index, + "seed": self._sampler._seed, + "config": copy.deepcopy(config), + } + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore without sampling or decoding, even before trainer setup. + + Data requirements and the ragged layout may still be registered after + this call. Their identity and the saved offset are checked lazily on + the first batch, before decoding it. Rebuilding only the saved pass + skips batch indices, never the frames of earlier batches or epochs. + """ + if self._closed: + raise RuntimeError("cannot restore a closed LMDB batch iterator") + if state.get("version") != 1: + raise ValueError("Unsupported LMDB data progress version") + for key in ("epoch", "batch_index", "seed"): + if type(state.get(key)) is not int or state[key] < 0: + raise ValueError(f"LMDB data progress {key} must be a non-negative int") + config = state.get("config") + if not isinstance(config, dict): + raise ValueError("LMDB data progress is missing its sampling configuration") + self._check_restore_config(config, final=False) + self._discard_prefetch() + self._sampler.set_seed(state["seed"]) + self._sampler.set_epoch(state["epoch"]) + self._epoch = self._next_epoch = state["epoch"] + self._next_batch_index = self._resume_batch_index = state["batch_index"] + self._restore_config = copy.deepcopy(config) + self._queued_pass_length = 0 + self._queued_nframes = 0.0 + self._last_pass_length = None + self._iterator = self._iter_epoch() def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]: """Submit one batch as balanced contiguous chunks. @@ -1897,7 +2094,7 @@ def _schedule(self, indices: list[int]) -> None: def batch_pass_length(self) -> int: """Length of the last returned batch's pass, or the first pending pass.""" if self._last_pass_length is None: - return len(self._sampler) + return self._queued_pass_length or len(self._sampler) return self._last_pass_length @property @@ -1910,11 +2107,8 @@ def closed(self) -> bool: """Whether this iterator has released its resources.""" return self._closed - def close(self) -> None: - """Finish or cancel prefetched work and release the decoder pool.""" - if self._closed: - return - self._closed = True + def _discard_prefetch(self) -> None: + """Discard queued reads while preserving the shared pool's health state.""" if self._pending is not None: running = [ future for future in self._pending.futures if not future.cancel() @@ -1923,16 +2117,18 @@ def close(self) -> None: self._lose_the_pool(running) self._pending = None self._deferred_indices = None + + def close(self) -> None: + """Finish or cancel prefetched work and release the decoder pool.""" + if self._closed: + return + self._closed = True + self._discard_prefetch() if self._pool is not None: self._pool = None _release_lmdb_executor(self._num_workers) -def is_lmdb(systems: str) -> bool: - """Check if systems points to an LMDB dataset.""" - return systems.endswith(".lmdb") or Path(systems, "data.mdb").is_file() - - def _parse_metadata(meta: dict) -> tuple[int, str, list[int]]: """Parse LMDB metadata into (nframes, frame_fmt, natoms_per_type). @@ -2055,6 +2251,8 @@ class LmdbDataReader: Reads LMDB frames and returns dicts of numpy arrays. Backend-specific Dataset classes (PyTorch, JAX, etc.) wrap this. + Element compatibility is checked on decoded real atoms, not on unused + entries in the dataset's element table. An LMDB typically holds frames of many different atom counts. The ``batch_size`` rule decides how those frames are grouped: @@ -2121,25 +2319,8 @@ def __init__( self._natoms = sum(self._natoms_per_type) self._ntypes = len(type_map) - # Build type remapping if LMDB's type_map differs from model's type_map - lmdb_type_map = meta.get("type_map") - self._lmdb_type_map = lmdb_type_map - self._type_remap: np.ndarray | None = None - if lmdb_type_map is not None and list(lmdb_type_map) != list(type_map): - # Build remap: lmdb_type_idx -> model_type_idx - remap = np.empty(len(lmdb_type_map), dtype=np.int32) - for i, name in enumerate(lmdb_type_map): - if name not in type_map: - raise ValueError( - f"Element '{name}' in LMDB type_map {lmdb_type_map} " - f"not found in model type_map {type_map}" - ) - remap[i] = type_map.index(name) - self._type_remap = remap - log.info( - f"Type remapping: LMDB {lmdb_type_map} -> model {type_map}, " - f"remap={remap}" - ) + self._lmdb_type_map = meta.get("type_map") + self._type_remap = _make_type_remap(self._lmdb_type_map, type_map) # The parent transaction serves synchronous reads. Decoder workers open # independent process-local environments and transactions. @@ -2234,7 +2415,7 @@ def __init__( # Re-key _frame_nlocs / _frame_system_ids into the dataset-index # space so that every downstream consumer (nloc_groups, system_groups, - # LmdbBatchSampler, _expand_indices_by_blocks) operates in a single, + # LmdbBatchSampler, _resample_indices_by_blocks) operates in a single, # self-consistent indexing scheme. keys_are_identity = retained_keys.size == self.nframes self._keys_are_identity = keys_are_identity @@ -2271,6 +2452,7 @@ def __init__( ).tolist() else: self._system_nframes = [self.nframes] + self.sampling_nsystems = sum(count > 0 for count in self._system_nframes) # Nominal batch size, reported to callers that want a single number. # The sampler never uses it: same-nloc modes go through @@ -2303,6 +2485,7 @@ def __init__( data_requirements=self._data_requirements, dataset=self.lmdb_path, type_map=self._type_map, + lmdb_type_map=self._lmdb_type_map, ) # Which fields carry an atom axis follows from the registered # requirements, so this cache is invalidated when they change. @@ -2457,7 +2640,7 @@ def batch_layout( # workers a chunk of it is shipped to. return BatchLayout.over( self._frame_nlocs[np.asarray(indices, dtype=np.int64)].astype(np.int64), - self.per_atom_strides(), + self.per_atom_strides(indices[0]), ragged=self._ragged_batches if ragged is None else ragged, ) @@ -2490,12 +2673,18 @@ def use_ragged_batches(self, ragged: bool) -> None: """ self._ragged_batches = ragged and self.mixed_nloc - def per_atom_strides(self) -> dict[str, int]: + def per_atom_strides(self, index: int = 0) -> dict[str, int]: """Return the leading-axis entries per atom of each per-atom field. Every frame of one LMDB exposes the same fields, so the classification is a property of the dataset and its registered requirements rather - than of a batch, and is resolved from the first frame once. + than of a batch, and is resolved from a selected frame once. + + Parameters + ---------- + index : int, optional + Dataset index of the frame used to initialize the cached strides. + Batch consumers supply a member of their selected indices. Returns ------- @@ -2503,7 +2692,7 @@ def per_atom_strides(self) -> dict[str, int]: Entries per atom, keyed by field name. """ if self._per_atom_strides is None: - frame = self[0] + frame = self[index] self._per_atom_strides = per_atom_strides( frame, resolve_per_atom_keys(frame, self._decode_config) ) @@ -2996,9 +3185,11 @@ def compute_block_targets( ) -> list[tuple[list[int], int]]: """Compute target frame count per block from auto_prob config. - Uses the same ``prob_sys_size_ext`` logic as the npy pipeline to parse - the ``auto_prob`` string, then converts per-system probabilities into - per-block target frame counts using the "max(frames/prob)" strategy. + Uniform sampling assigns each retained system the rounded-up mean frame + count. With ``N`` retained frames and ``M`` nonempty systems, the sampled + pass contains ``M * ceil(N / M)`` frames, fewer than ``N + M``. + Custom block rules use ``prob_sys_size_ext`` and the non-shrinking + ``max(frames / probability)`` strategy. Parameters ---------- @@ -3013,12 +3204,21 @@ def compute_block_targets( ------- list[tuple[list[int], int]] Each element is ``(system_indices_in_block, target_frame_count)``. - Returns empty list if no expansion is needed (all targets == actual). + Returns an empty list if no resampling is needed (all targets == actual). """ - from deepmd.utils.data_system import ( + from deepmd.utils.data_sampling import ( prob_sys_size_ext, ) + if auto_prob_style == "prob_uniform": + retained = [index for index, count in enumerate(system_nframes) if count > 0] + if not retained: + return [] + target = (sum(system_nframes) + len(retained) - 1) // len(retained) + if all(system_nframes[index] == target for index in retained): + return [] + return [([index], target) for index in retained] + # Parse block definitions from the auto_prob string # Format: "prob_sys_size;stt:end:weight;stt:end:weight;..." block_str = auto_prob_style.split(";")[1:] @@ -3107,118 +3307,76 @@ def compute_block_targets( return result -def _expand_indices_by_blocks( - indices: list[int], +def _resample_indices_by_blocks( + indices: np.ndarray | list[int], frame_system_ids: np.ndarray, - block_targets: list[tuple[list[int], int]], + lookup: np.ndarray, + group_targets: list[int], rng: np.random.Generator, - _block_total_actual: list[int] | None = None, - _sid_to_blk_arr: np.ndarray | None = None, - _group_block_targets: list[int] | None = None, ) -> list[int]: - """Expand frame indices according to block targets. + """Resample a homogeneous group's frames to its exact block targets. - For each block, computes the proportional target for the subset of - indices belonging to that block (within the current nloc group), - then applies full-copy + remainder sampling. + Targets are allocated globally by :func:`_allocate_group_block_targets`. + A smaller target selects frames without replacement; a larger target + uses full-copy + remainder sampling. Parameters ---------- - indices : list[int] - Frame indices in the current nloc group. + indices : numpy.ndarray or list[int] + Frame indices in the current homogeneous group. frame_system_ids : np.ndarray Per-frame system id for the entire dataset. - block_targets : list[tuple[list[int], int]] - Per-block (system_ids, total_target_frames). + lookup : numpy.ndarray + System-id to block-index table from :func:`system_block_lookup`. + group_targets : list[int] + Exact target for each block in this group. rng : np.random.Generator - RNG for remainder sampling. - _block_total_actual : list[int] or None - Pre-computed total actual frame count per block (across all nloc - groups). When provided, avoids an O(N) scan of frame_system_ids. - _sid_to_blk_arr : np.ndarray or None - Pre-computed lookup from :func:`system_block_lookup`. When provided, - avoids rebuilding the mapping for each call. - _group_block_targets : list[int] or None - Exact target for each block in this group. Production samplers - allocate these targets globally across all groups so independent - rounding cannot change a block's total size. + RNG for subset and remainder sampling. Returns ------- list[int] - Expanded indices. + Resampled indices. """ - n_blocks = len(block_targets) - - if _sid_to_blk_arr is None: - _sid_to_blk_arr = system_block_lookup(block_targets) - - sid_arr = np.asarray(frame_system_ids) idx_arr = np.asarray(indices, dtype=np.int64) - idx_blks = resolve_frame_blocks(idx_arr, sid_arr, _sid_to_blk_arr) - - # Pre-compute block_total_actual if not provided - if _block_total_actual is None and _group_block_targets is None: - _block_total_actual = count_group_blocks( - np.arange(sid_arr.size, dtype=np.int64), - sid_arr, - _sid_to_blk_arr, - n_blocks, - ).tolist() - - expanded_parts: list[np.ndarray] = [] + idx_blks = resolve_frame_blocks(idx_arr, frame_system_ids, lookup) + sampled_parts: list[np.ndarray] = [] # Unassigned indices unassigned_mask = idx_blks == -1 if np.any(unassigned_mask): - expanded_parts.append(idx_arr[unassigned_mask]) + sampled_parts.append(idx_arr[unassigned_mask]) - for blk_idx in range(n_blocks): + for blk_idx, target in enumerate(group_targets): blk_mask = idx_blks == blk_idx blk_idxs = idx_arr[blk_mask] n_actual = len(blk_idxs) if n_actual == 0: continue - if _group_block_targets is not None: - target = _group_block_targets[blk_idx] - else: - _, block_total_target = block_targets[blk_idx] - block_total_act = _block_total_actual[blk_idx] - - # Backward-compatible fallback for direct callers. Samplers pass - # exact group targets allocated by _allocate_group_block_targets. - if block_total_act > 0: - target = round(block_total_target * n_actual / block_total_act) - else: - target = n_actual - target = max(target, n_actual) # never shrink - if target < n_actual: - raise ValueError( - "Per-group auto-probability target cannot shrink original frames: " - f"target={target}, actual={n_actual}, block={blk_idx}" - ) + sampled_parts.append(rng.choice(blk_idxs, size=target, replace=False)) + continue # Full copies + remainder deficit = target - n_actual if deficit <= 0: - expanded_parts.append(blk_idxs) + sampled_parts.append(blk_idxs) else: full_copies = deficit // n_actual remainder = deficit % n_actual # Original + full copies if full_copies > 0: - expanded_parts.append(np.tile(blk_idxs, 1 + full_copies)) + sampled_parts.append(np.tile(blk_idxs, 1 + full_copies)) else: - expanded_parts.append(blk_idxs) + sampled_parts.append(blk_idxs) # Remainder: sample without replacement if remainder > 0: sampled = rng.choice(blk_idxs, size=remainder, replace=False) - expanded_parts.append(sampled) + sampled_parts.append(sampled) - if expanded_parts: - return np.concatenate(expanded_parts).tolist() + if sampled_parts: + return np.concatenate(sampled_parts).tolist() return [] @@ -3365,13 +3523,16 @@ def _allocate_group_block_targets( groups: list[np.ndarray], frame_system_ids: np.ndarray, block_targets: list[tuple[list[int], int]], + rng: np.random.Generator, ) -> list[list[int]]: """Allocate every block target exactly across homogeneous groups. - Original frames form the non-shrinking baseline. Each block's expansion + For expansion, original frames form the non-shrinking baseline. Each block's deficit is apportioned by actual group size using the integer largest- remainder method. Stable group order breaks equal-remainder ties, which keeps distributed ranks deterministic without floating-point rounding. + A smaller target instead samples group counts from a uniform subset of + the block's frames, so no small group is permanently excluded by rounding. """ lookup = system_block_lookup(block_targets) group_actual = [ @@ -3384,11 +3545,6 @@ def _allocate_group_block_targets( group_targets = [counts.copy() for counts in group_actual] for block_index, (_system_ids, block_target) in enumerate(block_targets): block_actual = sum(counts[block_index] for counts in group_actual) - if block_target < block_actual: - raise ValueError( - "Auto-probability block target cannot shrink original frames: " - f"target={block_target}, actual={block_actual}, block={block_index}" - ) if block_actual == 0: if block_target != 0: raise ValueError( @@ -3397,6 +3553,20 @@ def _allocate_group_block_targets( ) continue + if block_target < block_actual: + # Uniform positions in the concatenated groups give the group + # counts of a global sample without replacement. Sampling within + # each group then gives every block frame the same inclusion chance. + stops = np.cumsum( + [counts[block_index] for counts in group_actual], dtype=np.int64 + ) + positions = rng.choice(block_actual, size=block_target, replace=False) + sampled_groups = np.searchsorted(stops, positions, side="right") + counts = np.bincount(sampled_groups, minlength=len(groups)) + for group_index, target in enumerate(counts): + group_targets[group_index][block_index] = int(target) + continue + deficit = block_target - block_actual remainders: list[tuple[int, int]] = [] allocated = 0 @@ -3541,7 +3711,7 @@ def _build_all_batches( Random number generator (deterministic for reproducibility). block_targets : list[tuple[list[int], int]] or None Per-block (system_ids, target_frame_count) from compute_block_targets. - When provided, indices are expanded via full-copy + remainder sampling. + When provided, each block is resampled to its target frame count. Returns ------- @@ -3561,22 +3731,21 @@ def _build_all_batches( if block_targets and reader.frame_system_ids is not None: sid_arr = reader.frame_system_ids group_block_targets = _allocate_group_block_targets( - groups, sid_arr, block_targets + groups, sid_arr, block_targets, rng ) sid_to_blk_arr = system_block_lookup(block_targets) for group_index, original_indices in enumerate(groups): indices = original_indices - # Expand each group independently using targets that were allocated + # Resample each group independently using targets that were allocated # globally, so that a block's total size is preserved exactly. if block_targets and sid_arr is not None and group_block_targets is not None: - indices = _expand_indices_by_blocks( + indices = _resample_indices_by_blocks( indices, sid_arr, - block_targets, + sid_to_blk_arr, + group_block_targets[group_index], rng, - _sid_to_blk_arr=sid_to_blk_arr, - _group_block_targets=group_block_targets[group_index], ) if shuffle: # ``permutation`` returns a copy, which matters because a group @@ -3630,7 +3799,7 @@ class LmdbBatchSampler: seed : int or None Random seed for reproducibility. block_targets : list[tuple[list[int], int]] or None - Per-block expansion targets from compute_block_targets. + Per-block resampling targets from compute_block_targets. """ def __init__( @@ -3642,11 +3811,16 @@ def __init__( ) -> None: self._reader = reader self._shuffle = shuffle - self._seed = seed + self._configured_seed = seed + # Draw entropy once, not once per epoch, so even an unseeded run can + # reconstruct its next pass from a compact checkpoint. + self._seed = int(np.random.SeedSequence().entropy) if seed is None else seed self._epoch = 0 self._block_targets = block_targets self._batches: list[list[int]] | None = None + self._batch_frame_counts: np.ndarray | None = None self._data_requirements_revision = -1 + self._ragged_batches: bool | None = None def batches(self) -> list[list[int]]: """Return the batch list of the pending pass, building it if needed. @@ -3657,17 +3831,29 @@ def batches(self) -> list[list[int]]: Dataset indices grouped into batches. """ revision = self._reader.data_requirements_revision - if self._batches is None or self._data_requirements_revision != revision: - seed = None if self._seed is None else self._seed + self._epoch + ragged = self._reader.mixed_nloc and self._reader.ragged_batches + if ( + self._batches is None + or self._data_requirements_revision != revision + or self._ragged_batches != ragged + ): self._batches = _build_all_batches( self._reader, self._shuffle, - np.random.default_rng(seed), + np.random.default_rng(self._seed + self._epoch), self._block_targets, ) self._data_requirements_revision = revision + self._ragged_batches = ragged + self._batch_frame_counts = None return self._batches + def set_seed(self, seed: int) -> None: + """Restore the effective seed, invalidating any cached pass.""" + if seed != self._seed: + self._seed = seed + self._batches = None + def set_epoch(self, epoch: int) -> None: """Select the pass to serve, discarding any pass still pending. @@ -3690,6 +3876,15 @@ def __len__(self) -> int: """Number of batches the pending pass holds.""" return len(self.batches()) + def batch_frame_counts(self) -> np.ndarray: + """Return the number of frames in each pending optimizer batch.""" + batches = self.batches() + if self._batch_frame_counts is None: + self._batch_frame_counts = np.fromiter( + (len(batch) for batch in batches), dtype=np.float64 + ) + return self._batch_frame_counts + @property def total_batches(self) -> int: """Number of batches the pending pass holds over the whole dataset. @@ -3727,7 +3922,7 @@ class DistributedLmdbBatchSampler: seed : int or None Base seed for deterministic RNG. All ranks must use the same seed. block_targets : list[tuple[list[int], int]] or None - Per-block expansion targets from compute_block_targets. + Per-block resampling targets from compute_block_targets. """ def __init__( @@ -3743,10 +3938,18 @@ def __init__( self._rank = rank self._world_size = world_size self._shuffle = shuffle + self._configured_seed = seed self._seed = seed if seed is not None else 0 self._epoch = 0 self._block_targets = block_targets self._global: LmdbBatchSampler | None = None + self._batch_frame_counts: np.ndarray | None = None + + def set_seed(self, seed: int) -> None: + """Restore the common seed, invalidating any cached global pass.""" + if seed != self._seed: + self._seed = seed + self._global = None def set_epoch(self, epoch: int) -> None: """Set epoch for deterministic cross-rank shuffling. @@ -3767,7 +3970,11 @@ def _global_batches(self) -> list[list[int]]: seed=self._seed + self._epoch, block_targets=self._block_targets, ) - return self._global.batches() + previous = self._global._batches + batches = self._global.batches() + if batches is not previous: + self._batch_frame_counts = None + return batches def __iter__(self) -> Iterator[list[int]]: """Yield this rank's partition of the global batch list.""" @@ -3804,6 +4011,26 @@ def __len__(self) -> int: """Number of batches for this rank.""" return len(self._partition_batches(self._global_batches())) + def batch_frame_counts(self) -> np.ndarray: + """Return rank-averaged frame counts for each pending optimizer step. + + Padding matches the strided partition, so every rank obtains identical + source-selection weights even when their local batch sizes differ. + Custom partitioning strategies must provide their corresponding + rank-averaged counts when used with uniform source sampling. + """ + self._global_batches() + if self._batch_frame_counts is None: + counts = self._global.batch_frame_counts() + self._batch_frame_counts = ( + np.resize(counts, len(self) * self._world_size) + .reshape(-1, self._world_size) + .mean(axis=1) + if counts.size + else counts + ) + return self._batch_frame_counts + @property def total_batches(self) -> int: """Number of batches one full pass holds, before the per-rank split.""" @@ -3889,6 +4116,8 @@ class LmdbTestData: For mixed-nloc datasets, frames are grouped by nloc. get_test(nloc=...) returns data for a specific group. + Only types used by decoded frames must be supported by the model; + unused element-table entries do not restrict the model's type map. """ def __init__( @@ -3912,28 +4141,8 @@ def __init__( len(self._type_map) if self._type_map else len(self._natoms_per_type) ) - # Build type remapping if LMDB's type_map differs from model's type_map - lmdb_type_map = meta.get("type_map") - self._lmdb_type_map = lmdb_type_map - self._type_remap: np.ndarray | None = None - if ( - lmdb_type_map is not None - and self._type_map - and list(lmdb_type_map) != list(self._type_map) - ): - remap = np.empty(len(lmdb_type_map), dtype=np.int32) - for i, name in enumerate(lmdb_type_map): - if name not in self._type_map: - raise ValueError( - f"Element '{name}' in LMDB type_map {lmdb_type_map} " - f"not found in model type_map {self._type_map}" - ) - remap[i] = self._type_map.index(name) - self._type_remap = remap - log.info( - f"LmdbTestData type remapping: LMDB {lmdb_type_map} -> " - f"model {self._type_map}, remap={remap}" - ) + self._lmdb_type_map = meta.get("type_map") + self._type_remap = _make_type_remap(self._lmdb_type_map, self._type_map or None) # Select the frames to test, without reading any of them. Atom counts # come from the metadata, so grouping, shuffling and truncation are @@ -3954,6 +4163,7 @@ def __init__( data_requirements=self._requirements, dataset=self.lmdb_path, type_map=self._type_map, + lmdb_type_map=self._lmdb_type_map, ) # Detect PBC from the first retained frame. @@ -3961,11 +4171,13 @@ def __init__( first = next( (indices[0] for indices in self._nloc_groups.values() if len(indices)), None ) - probe = self._read_frames([first]) if first is not None else [] - if probe: - box = probe[0].get("box") - if not isinstance(box, np.ndarray) or np.allclose(box, 0.0): - self.pbc = False + if first is not None: + with self._env.begin() as transaction: + raw = transaction.get(format(first, self._frame_fmt).encode()) + if raw is not None: + frame = _remap_keys(msgpack.unpackb(raw, raw=False)) + box = _decode_value(frame.get("box"), copy_arrays=False) + self.pbc = isinstance(box, np.ndarray) and not np.allclose(box, 0.0) self.mixed_type = True diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 6895ce3e1d..f13d80cb46 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -5,8 +5,12 @@ ``graph`` (the ``NeighborGraph``/``GraphLayout`` contract + derived node-validity + edge padding), ``csr`` (backend-agnostic CSR construction and canonicalization), ``builder`` (the carry-all ``build_neighbor_graph`` dispatcher + the -``from_dense_quartet`` legacy converter), ``segment`` (mask-aware -segment-reduction toolkit), and ``derivatives`` (edge force/virial assembly). +``from_dense_quartet`` legacy converter), ``pair_graph`` (the complete +intra-frame pair topology of interactions without a cutoff), +``reciprocal_graph`` (the reciprocal-space topologies of periodic frames), +``topology`` (building and laying out the topologies a model declares), +``segment`` (mask-aware segment-reduction toolkit), and ``derivatives`` (edge +and pair force/virial assembly). See the design discussion wanghan-iapcm/deepmd-kit#4. """ @@ -33,6 +37,7 @@ ) from .derivatives import ( edge_force_virial, + pair_force_virial, ) from .env import ( edge_env_mat, @@ -53,19 +58,56 @@ pad_and_guard_angles, pad_and_guard_edges, ) +from .pair_graph import ( + ALL_PAIRS, + PairGraph, + build_pair_graph, + check_complete_frames, + unordered_pair_index, +) from .pairs import ( center_edge_pairs, ) +from .reciprocal_graph import ( + K_BLOCK, + K_GRID, + PME_MESH, + KGrid, + MeshGrid, + build_fixed_k_grid, + build_k_grid, + build_mesh_grid, + check_periodic_cells, + enumerate_k_grid_bounds, + k_grid_bounds_of_cells, + mesh_dims_of_cells, +) from .segment import ( segment_max, segment_mean, segment_softmax, segment_sum, ) +from .topology import ( + TOPOLOGY_TYPES, + build_topologies, + cast_topology_geometry, + flatten_topologies, + topology_flat_fields, + unflatten_topologies, +) __all__ = [ + "ALL_PAIRS", + "K_BLOCK", + "K_GRID", + "PME_MESH", + "TOPOLOGY_TYPES", "GraphLayout", + "KGrid", + "MeshGrid", "NeighborGraph", + "PairGraph", "angle_padding_fraction", "angle_to_edge_sum", "angle_to_node_sum", @@ -75,25 +117,41 @@ "attach_edge_csr", "build_angle_index", "build_edge_csr", + "build_fixed_k_grid", + "build_k_grid", + "build_mesh_grid", "build_neighbor_graph", "build_neighbor_graph_ase", + "build_pair_graph", + "build_topologies", "canonicalize_neighbor_graph", + "cast_topology_geometry", "center_edge_pairs", + "check_complete_frames", + "check_periodic_cells", "compact_nodes", "edge_env_mat", "edge_force_virial", + "enumerate_k_grid_bounds", "expand_node_values", + "flatten_topologies", "frame_id_from_n_node", "from_dense_quartet", "graph_angle_cos", "graph_from_dense_quartet", + "k_grid_bounds_of_cells", + "mesh_dims_of_cells", "neighbor_graph_from_ijs", "node_ownership_mask", "node_validity_mask", "pad_and_guard_angles", "pad_and_guard_edges", + "pair_force_virial", "segment_max", "segment_mean", "segment_softmax", "segment_sum", + "topology_flat_fields", + "unflatten_topologies", + "unordered_pair_index", ] diff --git a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py index 09111fa5d5..7f643bd418 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -7,7 +7,9 @@ Conventions (see the unified edge-nlist design discussion, wanghan-iapcm/deepmd-kit#4): edge_vec_e = r_src - r_dst ; F_k = sum_{dst=k} g - sum_{src=k} g per-edge virial w_e = -g_e (x) edge_vec_e - atom virial attributed FULL-TO-src (canonical TF==pt-legacy convention) + atom virial attributed FULL-TO-src (canonical TF==pt-legacy convention); + a pair set (pair_force_virial) splits its pair virial evenly between the + two endpoints instead per-frame virial = sum over the edges of that frame of w_e (DeePMD virials are per frame; a multi-frame NeighborGraph must NOT collapse frames) Padding/guard edges (edge_mask == 0) are zeroed before any scatter. @@ -118,3 +120,71 @@ def edge_force_virial( node_frame = frame_id_from_n_node(n_node, n_total=n_out) # (N,) frame per node virial = segment_sum(atom_virial, node_frame, nf) # (nf, 3, 3) return force, atom_virial, virial + + +def pair_force_virial( + g_p: Array, + pair_vec: Array, + pair_index: Array, + pair_mask: Array, + n_node: Array, + node_capacity: int | None = None, +) -> tuple[Array, Array, Array]: + """Assemble per-node force/atom-virial and PER-FRAME virial of a pair set. + + A pair set lists every unordered pair once (see + :class:`~deepmd.dpmodel.utils.neighbor_graph.pair_graph.PairGraph`), so its + per-pair gradient ``g_p = dE/d(pair_vec)`` acts on both endpoints: + ``pair_vec = r_src - r_dst`` gives ``F_dst += g_p`` and ``F_src -= g_p``, + as for an edge. The pair virial ``w = -g_p (x) pair_vec`` is split evenly + between the two endpoints, which keeps the atomic virial independent of + the order in which a pair is listed; its per-frame sum is the sum of + ``w`` over the pairs of that frame. + + Parameters + ---------- + g_p + (P, 3) per-pair gradient ``dE/d(pair_vec)``. + pair_vec + (P, 3) per-pair displacement ``r_src - r_dst``. + pair_index + (2, P) ``[src, dst]`` node endpoints of each pair. + pair_mask + (P,) boolean real-pair mask; guard pairs (``False``) are zeroed + before any scatter. + n_node + (nf,) per-frame node counts. + node_capacity + Size of the node axis ``N``; ``None`` takes ``sum(n_node)``. + + Returns + ------- + force + (N, 3) per-node force. + atom_virial + (N, 3, 3) per-node virial, half of each pair virial per endpoint. + virial + (nf, 3, 3) per-frame virial. + """ + xp = array_api_compat.array_namespace(g_p) + n_out = node_capacity if node_capacity is not None else int(xp.sum(n_node)) + nf = n_node.shape[0] + if isinstance(n_out, int) and n_out == 0: + device = array_api_compat.device(g_p) + return ( + xp.zeros((0, 3), dtype=g_p.dtype, device=device), + xp.zeros((0, 3, 3), dtype=g_p.dtype, device=device), + xp.zeros((nf, 3, 3), dtype=g_p.dtype, device=device), + ) + g = g_p * xp.astype(pair_mask[:, None], g_p.dtype) + # As for edges, the modulo bounds the guard endpoints under export. + src = pair_index[0] % n_out + dst = pair_index[1] % n_out + force = segment_sum(g, dst, n_out) - segment_sum(g, src, n_out) + w_half = -0.5 * (g[:, :, None] * pair_vec[:, None, :]) # (P, 3, 3) + atom_virial = segment_sum(w_half, src, n_out) + segment_sum(w_half, dst, n_out) + # Both endpoints of a pair share a frame, so reducing the per-node virial + # over nodes equals reducing the pair virial over pairs. + node_frame = frame_id_from_n_node(n_node, n_total=n_out) + virial = segment_sum(atom_virial, node_frame, nf) + return force, atom_virial, virial diff --git a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py index d7e1a5274e..270bbadf3d 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py +++ b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py @@ -110,16 +110,19 @@ def neighbor_graph_from_ijs( j = xp.astype(xp.asarray(j, device=dev), xp.int64) nframe_id = xp.astype(xp.asarray(nframe_id, device=dev), xp.int64) # Within-frame indices become node indices through the frame offsets. - offset = xp.take(xp.cumulative_sum(n_node) - n_node, nframe_id, axis=0) - i_flat = i + offset - j_flat = j + offset + if nf == 1: + i_flat, j_flat = i, j + else: + offset = xp.take(xp.cumulative_sum(n_node) - n_node, nframe_id, axis=0) + i_flat = i + offset + j_flat = j + offset r_i = xp.take(coord_flat, i_flat, axis=0) r_j = xp.take(coord_flat, j_flat, axis=0) edge_vec = r_j - r_i if box is not None: box = xp.asarray(box, device=dev) box = xp.reshape(box, (nf, 3, 3)) - box_per_edge = xp.take(box, nframe_id, axis=0) # (E, 3, 3) + box_per_edge = box[0] if nf == 1 else xp.take(box, nframe_id, axis=0) S = xp.astype(xp.asarray(S, device=dev), box.dtype) # S @ box per edge via broadcast sum (NEVER np.einsum, which breaks on torch): # shift[e, b] = sum_a S[e, a] * box[e, a, b] diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py new file mode 100644 index 0000000000..aa0bd0ef39 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Complete intra-frame pair topology for interactions without a cutoff. + +A module whose interaction acts between every two atoms of a system, such as +a non-periodic long-range module, cannot read its geometry from the +short-range :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph`. It +declares the :data:`ALL_PAIRS` topology instead: every unordered pair of valid +nodes sharing a frame, listed once with ``src < dst``. The pairs are generated +from the per-frame node counts alone, so their storage is the true pair count +``sum_f n_f (n_f - 1) / 2`` with neither a distance search nor padding to the +widest frame. Pairs of types excluded by the model-level pair exclusion are +left out, as they are left out of the neighbor graph. Like the edge axis of +the neighbor graph, the pair axis ends in masked guard pairs, which keep its +length a dynamic dimension of its own however few real pairs a batch holds. + +Geometry enters only through :attr:`PairGraph.vec`, which an autograd backend +makes a leaf next to the short-range ``edge_vec``; its gradient is pulled back +to the nodes by :func:`~deepmd.dpmodel.utils.neighbor_graph.derivatives.pair_force_virial`. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + ClassVar, +) + +import array_api_compat + +from .graph import ( + frame_id_from_n_node, + pad_and_guard_edges, +) +from .segment import ( + segment_sum, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + +ALL_PAIRS = "all_pairs" +"""Name of the topology holding every unordered pair of valid nodes of a frame.""" + + +@dataclass +class PairGraph: + """Unordered pairs of nodes on the flat node axis of a NeighborGraph.""" + + index: Array + """(2, P) int64 [src, dst] with src < dst; both endpoints in one frame.""" + vec: Array + """(P, 3) float r_src - r_dst; the geometry leaf of the pair interaction.""" + mask: Array + """(P,) bool real (1) vs guard (0) pairs; guards join node 0 to itself.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ("vec",) + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("vec",) + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "index": {1: "entry"}, + "vec": {0: "entry"}, + "mask": {0: "entry"}, + } + """Dynamic axes of every field: the pair axis.""" + + +def unordered_pair_index(n_node: Array, node_mask: Array | None = None) -> Array: + """Enumerate the unordered pairs of valid nodes within every frame. + + Within a frame of ``n`` valid nodes the pairs follow the row-major order + of the strict upper triangle, ``(0, 1), (0, 2), ..., (0, n-1), (1, 2), ...``. + Row ``i`` starts at rank ``S(i) = i (2n - 1 - i) / 2``; the row of rank + ``k`` is the floor of the smaller root of ``S(i) = k``, evaluated in + float64 and corrected by one integer step against rounding. + + The pair count is data dependent, so the enumeration runs eagerly. + + Parameters + ---------- + n_node + Per-frame node counts with shape ``(nf,)``. + node_mask + Validity of the flat nodes with shape ``(N,)``; ``None`` takes every + node as valid. Invalid nodes appear in no pair. + + Returns + ------- + Array + The ``[src, dst]`` node indices with shape ``(2, P)``, int64. + """ + xp = array_api_compat.array_namespace(n_node) + device = array_api_compat.device(n_node) + n_node = xp.astype(n_node, xp.int64) + if node_mask is None: + count = n_node + else: + frame_id = frame_id_from_n_node(n_node, n_total=node_mask.shape[0]) + count = segment_sum(xp.astype(node_mask, xp.int64), frame_id, n_node.shape[0]) + n_pair = count * (count - 1) // 2 + total = int(xp.sum(n_pair)) + if total == 0: + return xp.zeros((2, 0), dtype=xp.int64, device=device) + pair_end = xp.cumulative_sum(n_pair) + rank = xp.arange(total, dtype=xp.int64, device=device) + pair_frame = xp.astype(xp.searchsorted(pair_end, rank, side="right"), xp.int64) + local = rank - xp.take(pair_end - n_pair, pair_frame, axis=0) + n = xp.take(count, pair_frame, axis=0) + + def row_start(row: Array) -> Array: + return row * (2 * n - 1 - row) // 2 + + b = xp.astype(2 * n - 1, xp.float64) + discriminant = b * b - 8.0 * xp.astype(local, xp.float64) + discriminant = xp.where( + discriminant > 0.0, discriminant, xp.zeros_like(discriminant) + ) + row = xp.astype(xp.floor((b - xp.sqrt(discriminant)) / 2.0), xp.int64) + row = xp.where(local < row_start(row), row - 1, row) + row = xp.where(local >= row_start(row + 1), row + 1, row) + col = row + 1 + (local - row_start(row)) + first = xp.take(xp.cumulative_sum(count) - count, pair_frame, axis=0) + src = first + row + dst = first + col + if node_mask is not None: + # The ranks above count valid nodes only; map them to node positions. + position = xp.reshape(xp.nonzero(node_mask)[0], (-1,)) + position = xp.astype(position, xp.int64) + src = xp.take(position, src, axis=0) + dst = xp.take(position, dst, axis=0) + return xp.stack([src, dst]) + + +def build_pair_graph( + coord: Array, + n_node: Array, + atype: Array, + node_mask: Array | None = None, + pair_excl: PairExcludeMask | None = None, +) -> PairGraph: + """Build the :data:`ALL_PAIRS` topology of flat node coordinates. + + Parameters + ---------- + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + A differentiable input keeps :attr:`PairGraph.vec` differentiable. + n_node + Per-frame node counts with shape ``(nf,)``. + atype + Node types with shape ``(N,)``; invalid nodes carry any type. + node_mask + Validity of the nodes with shape ``(N,)``, or ``None`` when every + node is valid. + pair_excl + The model-level pair exclusion, or ``None``; pairs of excluded types + are left out. + + Returns + ------- + PairGraph + The unordered pairs of valid nodes and their displacements, followed + by two masked guard pairs. + """ + xp = array_api_compat.array_namespace(coord) + index = unordered_pair_index(n_node, node_mask) + if pair_excl is not None: + keep = pair_excl.build_edge_exclude_mask(index, atype) + index = xp.take(index, xp.reshape(xp.nonzero(keep)[0], (-1,)), axis=1) + vec = xp.take(coord, index[0], axis=0) - xp.take(coord, index[1], axis=0) + index, vec, mask = pad_and_guard_edges(index, vec, None) + return PairGraph(index=index, vec=vec, mask=mask) + + +# Guard pairs of the pair topology of a batch without non-periodic frames. +_GUARD_PAIRS = 4 + + +def empty_pair_graph(like: Array) -> PairGraph: + """The pair topology of a batch without non-periodic frames: guard pairs only. + + The guard pairs number :data:`_GUARD_PAIRS`, which differs from the two + rows of the index and the three components of the vectors: a compiled + graph would identify a pair axis of the length of one of those axes with + it, and split at the reductions over the pairs. + """ + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + index = xp.zeros((2, 0), dtype=xp.int64, device=device) + vec = xp.zeros((0, 3), dtype=like.dtype, device=device) + index, vec, mask = pad_and_guard_edges(index, vec, None, min_edges=_GUARD_PAIRS) + return PairGraph(index=index, vec=vec, mask=mask) + + +def check_complete_frames(box: Array | None) -> None: + """Reject a cell, since the all-pairs topology covers non-periodic frames. + + Parameters + ---------- + box + Simulation cells, or ``None`` for non-periodic systems. + + Raises + ------ + NotImplementedError + If a cell is given. + """ + if box is not None: + raise NotImplementedError( + "the all-pairs topology of the long-range modules describes " + "complete non-periodic systems only; a module evaluates periodic " + "cells with `periodic: true` (pass box=None for a non-periodic " + "system)" + ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py new file mode 100644 index 0000000000..c337cf40f4 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Reciprocal-space topologies of the periodic frames of a batch. + +A long-range module that serves periodic systems evaluates the lattice sum of +its pair kernel in reciprocal space, which needs the cell of every frame and +the absolute node positions rather than pair displacements. It declares one of +two topologies beside +:data:`~deepmd.dpmodel.utils.neighbor_graph.pair_graph.ALL_PAIRS`: + +- :data:`K_GRID` lists the wave numbers of every frame explicitly, ragged over + the frames, for a direct sum over them; +- :data:`PME_MESH` carries the cells and the mode indices of a real-space mesh + shared by the frames, for a smooth particle-mesh evaluation. + +Both carry the node positions and a zero strain as geometry: an autograd +backend makes them leaves, whose gradients are the force and the negative +virial of the reciprocal-space energy. A batch without cells comes back with +masked topologies, so a module serving periodic systems also serves +non-periodic batches through its pair topology. +""" + +from __future__ import ( + annotations, +) + +import math +import warnings +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + ClassVar, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + +K_GRID = "k_grid" +"""Name of the topology holding the reciprocal lattice points of periodic frames.""" + +PME_MESH = "pme_mesh" +"""Name of the topology carrying the cells and the mesh of periodic frames.""" + + +@dataclass +class KGrid: + r"""Reciprocal lattice points of the periodic frames of a batch, frame by frame. + + Row ``f`` of :attr:`vec` holds the wave numbers :math:`\mathbf{k} = 2\pi + B_f^{-1}\mathbf{n}` of frame ``f``, for integer :math:`\mathbf{n}`, built + eagerly outside any compiled region, in blocks of :data:`K_BLOCK` points + padded with masked guard points to the blocks of the longest row; row + ``f`` of :attr:`slot` holds the nodes of frame ``f``, padded with ``-1``. + A frame thus pairs its own wave numbers with its own nodes alone, at a + cost of the frame count times the padded wave-number count and the + padded node count, and every sum over the wave numbers of a frame runs + over the points of a block first. The block, point and slot axes are dynamic dimensions of + their own, of at least two entries. + + The energy is evaluated on the strained geometry of :attr:`strain`: + positions by ``S = I + strain``, wave numbers by :math:`S^{-T}` and volumes + by :math:`\det S`, so that the gradient with respect to the strain is the + negative virial; the phases :math:`\mathbf{k} \cdot \mathbf{r}` do not + change under the strain. The strain is zero, and the evaluation takes it + to first order, which fixes the energy and that gradient. The graph lower + makes :attr:`coord` and :attr:`strain` autograd leaves at its inputs. + """ + + vec: Array + """(nf, nb, kb, 3) float the wave numbers of every frame in blocks, in Å⁻¹.""" + mask: Array + """(nf, nb, kb) bool real (1) vs guard (0) points.""" + slot: Array + """(nf, ns) int64 the nodes of every frame, ``-1`` past its last node.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array + """(N, 3) float node positions, in Å.""" + strain: Array + """(nf, 3, 3) float zero strain of every frame.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ("vec", "volume", "coord", "strain") + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("coord", "strain") + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "vec": {0: "frame", 1: "entry", 2: "k_point"}, + "mask": {0: "frame", 1: "entry", 2: "k_point"}, + "slot": {0: "frame", 1: "slot"}, + "volume": {0: "frame"}, + "coord": {0: "node"}, + "strain": {0: "frame"}, + } + """Dynamic axes of every field: the frame, block, point, slot and node axes.""" + + +@dataclass +class MeshGrid: + r"""Cells and mesh of the periodic frames of a batch. + + The frames share a real-space mesh of :math:`n_1 \times n_2 \times n_3` + points spanning each cell, whose integer mode indices along every axis + are the fields :attr:`mode_1`, :attr:`mode_2` and :attr:`mode_3`: their + lengths are the mesh dimensions, so a compiled or exported graph takes the + mesh as dynamic dimensions. Every dimension is even. A batch without a + cell comes back with identity cells, a minimal mesh and an all-false + :attr:`mask`. + + The geometry contract is that of :class:`KGrid`: the energy is evaluated + with the positions strained by ``S = I + strain`` and the cells by + :math:`B S^{\mathsf{T}}`, the lattice vectors being the rows of + :math:`B`, so that the fractional coordinates are invariant and the + gradient with respect to the strain is the negative virial. The volumes + of the unstrained cells are computed eagerly, and the evaluation scales + them by :math:`\det S`, to first order in the strain. + """ + + box: Array + """(nf, 3, 3) float simulation cells; rows are the lattice vectors, in Å.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array + """(N, 3) float node positions, in Å.""" + strain: Array + """(nf, 3, 3) float zero strain of every frame.""" + mask: Array + """(nf,) bool periodic (1) vs non-periodic (0) frames.""" + mode_1: Array + """(n1,) float mode indices of the first mesh axis, in FFT order.""" + mode_2: Array + """(n2,) float mode indices of the second mesh axis, in FFT order.""" + mode_3: Array + """(n3,) float mode indices of the third mesh axis, in FFT order.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ( + "box", + "volume", + "coord", + "strain", + "mode_1", + "mode_2", + "mode_3", + ) + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("coord", "strain") + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "box": {0: "frame"}, + "volume": {0: "frame"}, + "coord": {0: "node"}, + "strain": {0: "frame"}, + "mask": {0: "frame"}, + "mode_1": {0: "mesh_1"}, + "mode_2": {0: "mesh_2"}, + "mode_3": {0: "mesh_3"}, + } + """Dynamic axes of every field: the frame axis, the node axis, the mesh axes.""" + + +def check_periodic_cells(box: Array) -> None: + """Reject a cell that spans no lattice. + + A cell whose volume is a vanishing fraction of the product of the lengths + of its lattice vectors, or that is not finite, spans no lattice; an + all-zero cell is no marker of a non-periodic frame, which takes no cell. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 9)`` or ``(nf, 3, 3)``; rows are + the lattice vectors. + + Raises + ------ + ValueError + If a cell spans no lattice. + """ + cells = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + volume = np.abs(np.linalg.det(cells)) + spanned = volume > 1e-6 * np.prod(np.linalg.norm(cells, axis=-1), axis=-1) + if not np.all(spanned): + f = int(np.argmin(spanned)) + raise ValueError( + f"the cell of periodic frame {f} has no volume; a periodic frame " + "needs three independent lattice vectors, and a non-periodic " + "system takes no cell" + ) + + +def _cells(box: Array) -> Array: + """The cells with shape ``(nf, 3, 3)``, from a ``(nf, 9)`` or ``(nf, 3, 3)`` array.""" + xp = array_api_compat.array_namespace(box) + return xp.reshape(box, (box.shape[0], 3, 3)) + + +def _volume(box: Array) -> Array: + """The cell volumes |a_1 . (a_2 x a_3)| with shape ``(nf,)``.""" + xp = array_api_compat.array_namespace(box) + triple = box[:, 0, :] * xp.linalg.cross(box[:, 1, :], box[:, 2, :]) + return xp.abs(xp.sum(triple, axis=-1)) + + +def _zero_strain(box: Array) -> Array: + xp = array_api_compat.array_namespace(box) + return xp.zeros( + (box.shape[0], 3, 3), dtype=box.dtype, device=array_api_compat.device(box) + ) + + +def k_grid_bounds_of_cells(box: Array, k_max: float) -> np.ndarray: + r"""Per-axis integer wave-number bounds the cells need at ``k_max``. + + A wave number :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` with + :math:`|\mathbf{k}| \le k_{max}` has :math:`|n_d| \le k_{max} |a_d| / 2\pi` + along every lattice vector :math:`\mathbf{a}_d` (Cauchy-Schwarz), so the + bound of every frame is the ceiling of the right-hand side. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + k_max + The wave-number cutoff, in Å⁻¹. + + Returns + ------- + np.ndarray + The per-frame, per-axis integer bounds with shape ``(nf, 3)``. + """ + cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + length = np.linalg.norm(cell, axis=-1) + return np.ceil(k_max * length / (2.0 * math.pi)).astype(np.int64) + + +def enumerate_k_grid_bounds(bounds: Array | list[int]) -> np.ndarray: + r"""Enumerate the integer wave-number indices within per-axis bounds. + + The grid holds every :math:`\mathbf{n} \neq \mathbf{0}` with + :math:`|n_d| \le n_d^{max}`, in lexicographic order with the first axis + outermost. The zero vector is left out: the transform of a Coulomb-like + kernel is singular there, and the tin-foil convention drops that mode. + + Parameters + ---------- + bounds + The per-axis bounds :math:`n_d^{max}` with shape ``(3,)``. + + Returns + ------- + np.ndarray + The integer indices with shape ``(nk, 3)``, int64. + + Raises + ------ + ValueError + If the bounds are not three non-negative integers. + """ + bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(-1) + if bounds.shape != (3,) or np.any(bounds < 0): + raise ValueError( + f"k-grid bounds must be three non-negative integers, got {bounds}" + ) + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=np.int64) for n in bounds], indexing="ij" + ) + n_grid = np.stack(grids, axis=-1).reshape(-1, 3) + return n_grid[np.any(n_grid != 0, axis=-1)] + + +K_BLOCK = 128 +"""Wave numbers per block of a row of a :class:`KGrid`.""" +# A guard point carries a nonzero wave number, whose kernel transform stays +# finite, and no weight. +_GUARD_K = (1e-6, 0.0, 0.0) +# The fewest entries of the block and slot axes, which keeps each a dynamic +# dimension of its own however few points or nodes a batch holds. +_MIN_ENTRIES = 2 + + +def _k_blocks(rows: list[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + """Lay the wave numbers of every frame out in blocks, padded with guard points. + + Returns the wave numbers with shape ``(nf, nb, K_BLOCK, 3)`` and their + mask with shape ``(nf, nb, K_BLOCK)``. + """ + longest = max([row.shape[0] for row in rows], default=0) + blocks = max(-(-longest // K_BLOCK), _MIN_ENTRIES) + vec = np.tile(np.asarray(_GUARD_K), (len(rows), blocks * K_BLOCK, 1)) + mask = np.zeros((len(rows), blocks * K_BLOCK), dtype=bool) + for f, row in enumerate(rows): + vec[f, : row.shape[0]] = row + mask[f, : row.shape[0]] = True + return ( + vec.reshape(len(rows), blocks, K_BLOCK, 3), + mask.reshape(len(rows), blocks, K_BLOCK), + ) + + +def _frame_slots(n_node: Array) -> np.ndarray: + """The nodes of every frame, one row each, padded with ``-1`` to the longest frame.""" + counts = np.asarray(to_numpy_array(n_node), dtype=np.int64).reshape(-1) + width = max(int(counts.max(initial=0)), _MIN_ENTRIES) + first = np.cumsum(counts) - counts + column = np.arange(width, dtype=np.int64)[None, :] + return np.where(column < counts[:, None], first[:, None] + column, -1) + + +def _k_grid( + box: Array, vec: np.ndarray, mask: np.ndarray, slot: np.ndarray, coord: Array +) -> KGrid: + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + return KGrid( + vec=xp.astype(xp.asarray(vec, device=device), box.dtype), + mask=xp.asarray(mask, device=device), + slot=xp.asarray(slot, device=device), + volume=xp.astype(_volume(box), box.dtype), + coord=coord, + strain=_zero_strain(box), + ) + + +def build_k_grid(box: Array, k_max: float, coord: Array, n_node: Array) -> KGrid: + r"""Enumerate the reciprocal lattice points of the periodic frames. + + With the lattice vectors as the rows of the cell :math:`B`, the wave + numbers of a frame are :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` over + integer :math:`\mathbf{n}` with :math:`0 < |\mathbf{k}| \le k_{max}`, so + that :math:`\mathbf{a}_d \cdot \mathbf{k} = 2\pi n_d`. Their count is data + dependent, so the enumeration runs eagerly. The points of every frame + follow the lexicographic order of their integer indices, and the rows of + the frames end in masked guard points up to the blocks of the longest + row. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + k_max + The wave-number cutoff, in Å⁻¹. + coord + Node positions with shape ``(N, 3)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + + Returns + ------- + KGrid + The reciprocal lattice points of every frame, one row each. + """ + box = _cells(box) + cells = np.asarray(to_numpy_array(box), dtype=np.float64) + bounds = k_grid_bounds_of_cells(cells, k_max) + rows = [] + for f, cell in enumerate(cells): + n_grid = enumerate_k_grid_bounds(bounds[f]).astype(np.float64) + k = n_grid @ (2.0 * math.pi * np.linalg.inv(cell)).T + rows.append(k[np.sum(k * k, axis=-1) <= k_max * k_max]) + vec, mask = _k_blocks(rows) + return _k_grid(box, vec, mask, _frame_slots(n_node), coord) + + +def build_fixed_k_grid( + box: Array, bounds: Array, k_max: float, coord: Array, n_node: Array +) -> KGrid: + r"""Evaluate a fixed integer k grid on the cells of the periodic frames. + + Every frame shares the integer indices within ``bounds`` (see + :func:`enumerate_k_grid_bounds`), so the blocks of a row have the static + count of the grid. The wave numbers of a frame are :math:`\mathbf{k} = + 2\pi B^{-1}\mathbf{n}`, and the points beyond the cutoff + :math:`|\mathbf{k}| \le k_{max}` and the guard points that complete the + last block are masked out. + + A frame whose cell needs indices beyond the grid, a larger cell than the + grid was sized for or a kernel whose cutoff has grown since, is evaluated + on the grid alone: its sum is truncated to the points of the grid, which + is issued as a warning. Its forces and virial remain the derivatives of + that energy. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + bounds + The per-axis bounds :math:`n_d^{max}` of the shared integer grid with + shape ``(3,)``. + k_max + The wave-number cutoff, in Å⁻¹. + coord + Node positions with shape ``(N, 3)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + + Returns + ------- + KGrid + The wave numbers of every frame over the shared grid, one row each. + """ + box = _cells(box) + cells = np.asarray(to_numpy_array(box), dtype=np.float64) + n_grid = enumerate_k_grid_bounds(bounds) + bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(3) + required = k_grid_bounds_of_cells(cells, k_max) + beyond = np.any(required > bounds[None, :], axis=1) + if np.any(beyond): + f = int(np.argmax(beyond)) + warnings.warn( + f"frame {f} needs integer k-grid bounds {required[f].tolist()}, " + f"beyond the fixed grid {bounds.tolist()}: its k-space sum is " + "truncated to the grid; enlarge `k_grid_bounds` for the full sum", + stacklevel=2, + ) + rec = 2.0 * math.pi * np.linalg.inv(cells) + k = np.einsum("jm,fdm->fjd", n_grid.astype(np.float64), rec) + vec, real = _k_blocks(list(k)) + mask = real & (np.sum(vec * vec, axis=-1) <= k_max * k_max) + return _k_grid(box, vec, mask, _frame_slots(n_node), coord) + + +def empty_k_grid(nf: int, coord: Array) -> KGrid: + """The k grid of a batch without periodic frames: guard points and no nodes.""" + xp = array_api_compat.array_namespace(coord) + device = array_api_compat.device(coord) + eye = xp.tile(xp.eye(3, dtype=coord.dtype, device=device)[None, ...], (nf, 1, 1)) + vec, mask = _k_blocks([np.zeros((0, 3), dtype=np.float64)] * nf) + return _k_grid( + eye, vec, mask, np.full((nf, _MIN_ENTRIES), -1, dtype=np.int64), coord + ) + + +def _even_fft_size(n: int) -> int: + """The smallest even integer of at least ``n`` with no prime factor above 7.""" + m = max(2, n + n % 2) + while True: + rest = m + for p in (2, 3, 5, 7): + while rest % p == 0: + rest //= p + if rest == 1: + return m + m += 2 + + +def mesh_dims_of_cells(box: Array, spacing: float) -> tuple[int, int, int]: + r"""The mesh that resolves every cell at ``spacing``. + + Along each lattice vector :math:`\mathbf{a}_d` the mesh holds the + smallest even FFT-friendly number of points (no prime factor above 7) of + at least :math:`\lceil \max_f |\mathbf{a}_d| / h \rceil`, the maximum + running over the frames, so that no frame is evaluated with a spacing + coarser than :math:`h`. The C++ interface sizes the mesh of a frozen + model by the same rule. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + spacing + The largest mesh spacing :math:`h`, in Å. + + Returns + ------- + tuple[int, int, int] + The mesh dimensions. + """ + cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + # The length is formed term by term, as the C++ interface forms it, so + # that both round the ceiling identically. + length = np.sqrt( + cell[..., 0] * cell[..., 0] + + cell[..., 1] * cell[..., 1] + + cell[..., 2] * cell[..., 2] + ) + longest = np.max(length, axis=0) + n1, n2, n3 = (_even_fft_size(math.ceil(float(a) / spacing)) for a in longest) + return n1, n2, n3 + + +def _mesh_modes(xp: object, n: int, dtype: object, device: object) -> Array: + """The integer mode indices of an axis of ``n`` points, in FFT order.""" + m = xp.arange(n, dtype=dtype, device=device) + return xp.where(m > n // 2, m - n, m) + + +def mesh_grid_of_dims( + box: Array, coord: Array, dims: tuple[int, int, int], mask: Array +) -> MeshGrid: + """Pack cells, node positions and a mesh of the given dimensions. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice + vectors. + coord + Node positions with shape ``(N, 3)``. + dims + The even mesh dimensions. + mask + Periodic (1) vs non-periodic (0) frames with shape ``(nf,)``. + + Returns + ------- + MeshGrid + The cells, the mesh and the geometry of the frames. + """ + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + mode_1, mode_2, mode_3 = (_mesh_modes(xp, n, box.dtype, device) for n in dims) + return MeshGrid( + box=box, + volume=xp.astype(_volume(box), box.dtype), + coord=coord, + strain=_zero_strain(box), + mask=mask, + mode_1=mode_1, + mode_2=mode_2, + mode_3=mode_3, + ) + + +def build_mesh_grid(box: Array, spacing: float, coord: Array) -> MeshGrid: + """Pack the cells of the periodic frames and their shared mesh. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + spacing + The largest mesh spacing, in Å (see :func:`mesh_dims_of_cells`). + coord + Node positions with shape ``(N, 3)``. + + Returns + ------- + MeshGrid + The cells, the mesh and the geometry of the frames, every frame + periodic. + """ + xp = array_api_compat.array_namespace(box) + box = _cells(box) + mask = xp.ones((box.shape[0],), dtype=xp.bool, device=array_api_compat.device(box)) + return mesh_grid_of_dims(box, coord, mesh_dims_of_cells(box, spacing), mask) + + +def empty_mesh_grid(nf: int, coord: Array) -> MeshGrid: + """The mesh grid of a batch without periodic frames: every frame masked.""" + xp = array_api_compat.array_namespace(coord) + device = array_api_compat.device(coord) + box = xp.tile(xp.eye(3, dtype=coord.dtype, device=device)[None, ...], (nf, 1, 1)) + mask = xp.zeros((nf,), dtype=xp.bool, device=device) + return mesh_grid_of_dims(box, coord, (2, 2, 2), mask) diff --git a/deepmd/dpmodel/utils/neighbor_graph/topology.py b/deepmd/dpmodel/utils/neighbor_graph/topology.py new file mode 100644 index 0000000000..e03e9aa2cf --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/topology.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Topologies an atomic model declares beside its neighbor graph. + +An interaction that the short-range neighbor graph cannot describe reads a +topology of its own, declared by name (see +:meth:`~deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel.required_topologies`): + +- :data:`~.pair_graph.ALL_PAIRS`, the pairs of the non-periodic frames; +- :data:`~.reciprocal_graph.K_GRID` and :data:`~.reciprocal_graph.PME_MESH`, + the reciprocal space of the periodic frames. + +Every topology is a dataclass of arrays that declares which of its fields +carry geometry (``GEOMETRY``), which an autograd backend makes leaves of the +energy (``LEAVES``), and the dynamic axes of every field (``AXES``), so that +casting, re-leafing and the inputs of a compiled or exported graph lower are +derived from the declaration alone. +""" + +from __future__ import ( + annotations, +) + +import dataclasses +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat + +from .pair_graph import ( + ALL_PAIRS, + PairGraph, + build_pair_graph, + check_complete_frames, + empty_pair_graph, +) +from .reciprocal_graph import ( + K_GRID, + PME_MESH, + KGrid, + MeshGrid, + build_fixed_k_grid, + build_k_grid, + build_mesh_grid, + check_periodic_cells, + empty_k_grid, + empty_mesh_grid, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterable, + Sequence, + ) + + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + Topology = PairGraph | KGrid | MeshGrid + +TOPOLOGY_TYPES: dict[str, type] = { + ALL_PAIRS: PairGraph, + K_GRID: KGrid, + PME_MESH: MeshGrid, +} +"""The dataclass of every topology name.""" + + +def _build_reciprocal( + name: str, + box: Array, + coord: Array, + n_node: Array, + params: dict[str, Any], +) -> KGrid | MeshGrid: + """Build a reciprocal-space topology of periodic frames from its parameters.""" + required = {K_GRID: "k_max", PME_MESH: "spacing"}[name] + if required not in params: + raise ValueError( + f"the {name!r} topology requires its `{required}` build parameter" + ) + check_periodic_cells(box) + if name == PME_MESH: + return build_mesh_grid(box, params["spacing"], coord) + if "bounds" in params: + return build_fixed_k_grid(box, params["bounds"], params["k_max"], coord, n_node) + return build_k_grid(box, params["k_max"], coord, n_node) + + +def build_topologies( + names: Iterable[str], + coord: Array, + n_node: Array, + box: Array | None = None, + node_mask: Array | None = None, + *, + atype: Array, + pair_excl: PairExcludeMask | None = None, + topology_params: dict[str, dict[str, Any]] | None = None, +) -> dict[str, Topology]: + """Build the topologies an atomic model declares, beside its neighbor graph. + + A batch is periodic when it carries cells. A declaration of the pairs + alone serves non-periodic batches only. Beside a reciprocal-space + topology, the pairs of a periodic batch come back guard-only and the + reciprocal-space topology of a non-periodic batch masked, so that one + model serves both. + + Parameters + ---------- + names + The declared topology names. + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + n_node + Per-frame node counts with shape ``(nf,)``. + box + Simulation cells with shape ``(nf, 9)`` or ``(nf, 3, 3)``, or + ``None`` for non-periodic systems. + node_mask + Validity of the nodes with shape ``(N,)``, or ``None`` when every + node is valid. + atype + Node types with shape ``(N,)``; invalid nodes carry any type. + pair_excl + The model-level pair exclusion, or ``None``. + topology_params + Build parameters by name: ``{"k_max": ...}`` for the k grid, with + ``"bounds"`` for a fixed integer grid (see + :func:`~.reciprocal_graph.build_fixed_k_grid`), and + ``{"spacing": ...}`` for the mesh. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + Each declared topology under its name; empty when none is declared. + + Raises + ------ + NotImplementedError + If cells are given for a declaration of the pairs alone, or a pair + exclusion for a periodic batch, whose reciprocal-space sum is global. + ValueError + If a name is unknown, a build parameter is missing, or a cell spans + no lattice. + """ + names = tuple(names) + unknown = [name for name in names if name not in TOPOLOGY_TYPES] + if unknown: + raise ValueError( + f"unknown topology {unknown[0]!r}; expected one of {tuple(TOPOLOGY_TYPES)}" + ) + periodic = any(name in (K_GRID, PME_MESH) for name in names) + if box is None: + nf = n_node.shape[0] + empty = {K_GRID: empty_k_grid, PME_MESH: empty_mesh_grid} + return { + name: ( + build_pair_graph(coord, n_node, atype, node_mask, pair_excl) + if name == ALL_PAIRS + else empty[name](nf, coord) + ) + for name in names + } + if not periodic: + if names: + check_complete_frames(box) + return {} + if pair_excl is not None: + raise NotImplementedError( + "the reciprocal-space sum of a periodic system is global; " + "`pair_exclude_types` cannot exclude pairs from it" + ) + params = topology_params or {} + return { + name: ( + empty_pair_graph(coord) + if name == ALL_PAIRS + else _build_reciprocal(name, box, coord, n_node, params.get(name, {})) + ) + for name in names + } + + +def topology_flat_fields(topology: Topology) -> tuple[Array, ...]: + """The fields of a topology in their declaration order.""" + return tuple( + getattr(topology, field.name) for field in dataclasses.fields(topology) + ) + + +def flatten_topologies(topologies: dict[str, Topology]) -> tuple[Array, ...]: + """Lay topologies out as positional inputs of a graph lower. + + A compiled or exported graph lower receives the fields of every topology, + in the order of the dict and of the declaration of each dataclass, after + its other inputs; :func:`unflatten_topologies` restores them. + + Parameters + ---------- + topologies + The topologies by name. + + Returns + ------- + tuple + The fields of every topology, concatenated. + """ + return tuple( + value + for topology in topologies.values() + for value in topology_flat_fields(topology) + ) + + +def unflatten_topologies( + names: Sequence[str], values: Sequence[Array] +) -> dict[str, Topology]: + """Restore topologies laid out by :func:`flatten_topologies`. + + Parameters + ---------- + names + The topology names, in the order they were laid out. + values + The fields of every topology, concatenated. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + The topologies by name. + + Raises + ------ + ValueError + If a name is unknown, or the number of values does not match the + layout of the names. + """ + topologies = {} + offset = 0 + for name in names: + if name not in TOPOLOGY_TYPES: + raise ValueError(f"unknown topology {name!r}") + cls = TOPOLOGY_TYPES[name] + arity = len(dataclasses.fields(cls)) + if offset + arity > len(values): + break + topologies[name] = cls(*values[offset : offset + arity]) + offset += arity + if len(topologies) != len(names) or offset != len(values): + raise ValueError( + f"{len(values)} flattened values given, which do not lay out the " + f"topologies {tuple(names)}" + ) + return topologies + + +def cast_topology_geometry(topology: Topology, dtype: Any) -> Topology: + """A copy of a topology with its geometry fields cast to ``dtype``.""" + xp = array_api_compat.array_namespace(getattr(topology, topology.GEOMETRY[0])) + return dataclasses.replace( + topology, + **{ + name: xp.astype(getattr(topology, name), dtype) + for name in topology.GEOMETRY + }, + ) diff --git a/deepmd/dpmodel/utils/neighbor_list.py b/deepmd/dpmodel/utils/neighbor_list.py index c9b6f5923b..bc8d379fba 100644 --- a/deepmd/dpmodel/utils/neighbor_list.py +++ b/deepmd/dpmodel/utils/neighbor_list.py @@ -63,8 +63,9 @@ class NeighborList: Implementations turn local coordinates into the extended representation: the coordinates and atom types of local-plus-ghost (periodic-image) atoms, a candidate neighbor list indexing the extended atoms, and a mapping from each - extended atom to its local owner. Implementations are stateless -- - ``rcut``/``sel`` are supplied by the model at call time. + extended atom to its local owner. Implementations may retain search storage + and neighbor candidates, while ``rcut``/``sel`` are supplied by the model at + call time. Returned geometry always describes the current coordinates. """ def build( diff --git a/deepmd/dpmodel/utils/training_utils.py b/deepmd/dpmodel/utils/training_utils.py index 6f46cb8c93..baa5a0dfb0 100644 --- a/deepmd/dpmodel/utils/training_utils.py +++ b/deepmd/dpmodel/utils/training_utils.py @@ -82,7 +82,8 @@ def resolve_model_prob( model_keys : list[str] List of model keys. model_prob_config : dict[str, float] | None - User-specified model probabilities. If None, use data size. + User-specified model probabilities. If omitted or empty, use native + data size: NumPy/HDF5 system counts or retained LMDB frame counts. model_training_data : dict[str, object] Training data for each model. rank : int, optional @@ -112,10 +113,12 @@ def resolve_model_prob( if is_node_main_process(rank): log.info( "training.model_prob is not set or empty; defaulting to the " - "number of systems per task." + "native data size per task (NumPy/HDF5 systems or LMDB frames)." ) for ii, model_key in enumerate(model_keys): - model_prob[ii] = float(_training_data_size(model_training_data[model_key])) + model_prob[ii] = float( + get_training_data_size(model_training_data[model_key]) + ) if not np.all(np.isfinite(model_prob)): raise ValueError("Model prob must be finite.") if np.any(model_prob < 0.0): @@ -126,7 +129,16 @@ def resolve_model_prob( return model_prob / sum_prob -def _training_data_size(training_data: object) -> int: +def get_training_data_size(training_data: object) -> int: + """Return native data size for default multi-task weights. + + A data adapter may expose a size distinct from its batch/statistics system + count. This preserves the PyTorch convention of counting NumPy/HDF5 + systems and retained LMDB frames, including within composed datasets. + """ + get_size = getattr(training_data, "get_model_sampling_size", None) + if get_size is not None: + return int(get_size()) if hasattr(training_data, "get_nsystems"): return int(training_data.get_nsystems()) if not isinstance(training_data, Sized): diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index d9972b967e..a6cdb3bfa9 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -34,7 +34,10 @@ from deepmd.utils.data import ( DeepmdData, ) -from deepmd.utils.data_system import ( +from deepmd.utils.data_config import ( + data_group_configs, +) +from deepmd.utils.data_path import ( process_systems, ) from deepmd.utils.weight_avg import ( @@ -98,34 +101,25 @@ def test( if numb_test == 0: # only float has inf, but should work for min numb_test = float("inf") - if train_json is not None: - jdata = j_loader(train_json) - jdata = update_deepmd_input(jdata) - data_params = jdata.get("training", {}).get("training_data", {}) - systems = data_params.get("systems") - if not systems: - raise RuntimeError("No training data found in input json") - root = Path(train_json).parent - if isinstance(systems, str): - systems = str((root / Path(systems)).resolve()) - else: - systems = [str((root / Path(ss)).resolve()) for ss in systems] - patterns = data_params.get("rglob_patterns", None) - all_sys = process_systems(systems, patterns=patterns) - elif valid_json is not None: - jdata = j_loader(valid_json) + if train_json is not None or valid_json is not None: + input_json = train_json if train_json is not None else valid_json + data_kind = "training" if train_json is not None else "validation" + jdata = j_loader(input_json) jdata = update_deepmd_input(jdata) - data_params = jdata.get("training", {}).get("validation_data", {}) - systems = data_params.get("systems") - if not systems: - raise RuntimeError("No validation data found in input json") - root = Path(valid_json).parent - if isinstance(systems, str): - systems = str((root / Path(systems)).resolve()) - else: - systems = [str((root / Path(ss)).resolve()) for ss in systems] - patterns = data_params.get("rglob_patterns", None) - all_sys = process_systems(systems, patterns=patterns) + data_params = jdata.get("training", {}).get(f"{data_kind}_data") + if not data_params: + raise RuntimeError(f"No {data_kind} data found in input json") + all_sys = [] + for group in data_group_configs(data_params): + if not group.get("systems"): + raise RuntimeError(f"No {data_kind} data found in input json") + all_sys.extend( + process_systems( + group["systems"], + patterns=group.get("rglob_patterns"), + root_dir=Path(input_json).parent, + ) + ) elif datafile is not None: with open(datafile) as datalist: all_sys = datalist.read().splitlines() diff --git a/deepmd/infer/deep_density.py b/deepmd/infer/deep_density.py new file mode 100644 index 0000000000..b1655c45b6 --- /dev/null +++ b/deepmd/infer/deep_density.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, +) + +from .deep_eval import ( + DeepEval, +) + + +class DeepDensity(DeepEval): + """Charge density evaluated on grid points. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + + @property + def output_def(self) -> ModelOutputDef: + """Get the output definition of this model. + + The output definition is identical to that of the density + fitting net of the model: the density is predicted on grid + points, and is neither reducible nor differentiable. + """ + return ModelOutputDef( + FittingOutputDef( + [ + OutputVariableDef( + "density", + [1], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + ) + + def eval( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, + grid: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + mixed_type: bool = False, + grid_batch_size: int | None = 4096, + **kwargs: Any, + ) -> np.ndarray: + """Evaluate the density on grid points. + + Parameters + ---------- + coords : np.ndarray + The coordinates of the atoms, in shape (nframes, natoms, 3). + cells : np.ndarray + The cell vectors of the system, in shape (nframes, 9). If the system + is not periodic, set it to None. + atom_types : list[int] or np.ndarray + The types of the atoms. If mixed_type is False, the shape is (natoms,); + otherwise, the shape is (nframes, natoms). + grid : np.ndarray + The coordinates of the grid points, in shape (nframes, ngrid, 3). + fparam : np.ndarray, optional + The frame parameters, by default None. + aparam : np.ndarray, optional + The atomic parameters, by default None. + mixed_type : bool, optional + Whether the atom_types is mixed type, by default False. + grid_batch_size : int or None, optional + Maximum number of grid points evaluated in one model call. Grid + chunks are concatenated in their original order. The default is + 4096; set to None to evaluate the full grid in one call. + **kwargs : dict[str, Any] + Keyword arguments. ``charge_spin`` (nframes, 2) is forwarded to + models with charge spin embedding; when omitted, such models use + their ``default_chg_spin`` if set. + + Returns + ------- + density + The density on the grid points, in shape (nframes, ngrid). + """ + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + _, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + grid = self._standard_grid(grid, nframes) + if grid_batch_size is not None and grid_batch_size <= 0: + raise ValueError("grid_batch_size must be positive or None") + if grid_batch_size is None or grid.shape[1] <= grid_batch_size: + grid_slices = (grid,) + else: + grid_slices = tuple( + grid[:, begin : begin + grid_batch_size] + for begin in range(0, grid.shape[1], grid_batch_size) + ) + density_chunks = [] + for grid_chunk in grid_slices: + results = self.deep_eval.eval( + coords, + cells, + atom_types, + False, + fparam=fparam, + aparam=aparam, + grid=grid_chunk, + **kwargs, + ) + density_chunks.append(results["density"].reshape(nframes, -1)) + return np.concatenate(density_chunks, axis=1) + + @staticmethod + def _standard_grid(grid: np.ndarray, nframes: int) -> np.ndarray: + """Normalise the grid coordinates to ``(nframes, ngrid, 3)``. + + The auto batcher slices every argument with ``ndim > 1`` along axis + 0, so a natural single-frame ``(ngrid, 3)`` input would silently be + truncated to one grid point; carry the frame dimension explicitly. + """ + arr = np.asarray(grid) + if arr.ndim == 2: + if nframes != 1: + raise ValueError( + f"grid of shape {arr.shape} is ambiguous for {nframes} " + "frames; pass grid with shape (nframes, ngrid, 3)" + ) + arr = arr[None, ...] + if arr.ndim != 3 or arr.shape[-1] != 3: + raise ValueError( + f"grid must have shape (nframes, ngrid, 3), got {arr.shape}" + ) + if arr.shape[0] != nframes: + raise ValueError(f"grid has {arr.shape[0]} frames but coord has {nframes}") + if arr.shape[1] == 0: + raise ValueError("grid must contain at least one point per frame") + return arr + + +__all__ = ["DeepDensity"] diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index ce23c5981b..96ef64a139 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -212,6 +212,7 @@ def eval( aparam=aparam, **kwargs, ) + energy = results["energy_redu"].reshape(nframes, 1) force = results["energy_derv_r"].reshape(nframes, natoms, 3) virial = results["energy_derv_c_redu"].reshape(nframes, 9) diff --git a/deepmd/infer/model_test/__init__.py b/deepmd/infer/model_test/__init__.py index 4afd626d0e..1152c3141d 100644 --- a/deepmd/infer/model_test/__init__.py +++ b/deepmd/infer/model_test/__init__.py @@ -14,6 +14,9 @@ Any, ) +from deepmd.infer.deep_density import ( + DeepDensity, +) from deepmd.infer.deep_dipole import ( DeepDipole, ) @@ -36,6 +39,9 @@ save_txt_file, test_chunk_atoms, ) +from deepmd.infer.model_test.density import ( + DensityTester, +) from deepmd.infer.model_test.dos import ( DosTester, ) @@ -56,6 +62,7 @@ __all__ = [ "ChunkContext", + "DensityTester", "DipoleTester", "DosTester", "EnerTester", @@ -96,6 +103,8 @@ def build_tester(dp: Any, *, atomic: bool) -> ModelTester: return tester(dp, atomic=atomic) if isinstance(dp, DeepDOS): return DosTester(dp, atomic=atomic) + if isinstance(dp, DeepDensity): + return DensityTester(dp, atomic=atomic) if isinstance(dp, DeepProperty): return PropertyTester(dp, atomic=atomic) if isinstance(dp, DeepGlobalPolar): diff --git a/deepmd/infer/model_test/density.py b/deepmd/infer/model_test/density.py new file mode 100644 index 0000000000..0ec2821b8a --- /dev/null +++ b/deepmd/infer/model_test/density.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of models predicting charge density on grid points.""" + +from deepmd.infer.model_test.base import ( + ChunkContext, + ModelTester, + _write_per_frame_details, +) +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.eval_metrics import ( + mae, + rmse, +) + +__all__ = ["DensityTester"] + + +class DensityTester(ModelTester): + """Test a model of charge density on grid points.""" + + report = ( + ("mae_density", "DENSITY MAE : {} units"), + ("rmse_density", "DENSITY RMSE : {} units"), + ) + + def add_data_requirements(self, data: DeepmdData) -> None: + """Declare the labels a density test consumes.""" + dp = self.dp + # The grid and the density are defined on grid points rather than on + # atoms, and their extent (ngrid) is not known until the data is + # loaded. They are declared frame-major so the loader keeps the + # (nframes, ngrid, ndof) layout without reshaping to natoms. + data.add( + "grid", + 3, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + data.add( + "density", + 1, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + if dp.get_dim_fparam() > 0: + data.add( + "fparam", + dp.get_dim_fparam(), + atomic=False, + must=not dp.has_default_fparam(), + high_prec=False, + ) + if dp.get_dim_aparam() > 0: + data.add( + "aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False + ) + if dp.has_chg_spin_ebd(): + data.add( + "charge_spin", + 2, + atomic=False, + must=not dp.has_default_chg_spin(), + high_prec=False, + ) + + def evaluate_chunk( + self, + data: DeepmdData, + test_data: dict, + context: ChunkContext, + ) -> dict[str, tuple[float, float]]: + """Evaluate one chunk of a density test.""" + dp = self.dp + mixed_type = data.mixed_type + nframes = test_data["box"].shape[0] + + coord = test_data["coord"].reshape([nframes, -1]) + box = test_data["box"] if data.pbc else None + if mixed_type: + atype = test_data["type"].reshape([nframes, -1]) + else: + atype = test_data["type"][0] + fparam = ( + test_data["fparam"] + if dp.get_dim_fparam() > 0 and test_data["find_fparam"] != 0.0 + else None + ) + aparam = test_data["aparam"] if dp.get_dim_aparam() > 0 else None + charge_spin = ( + test_data["charge_spin"] + if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 0.0) != 0.0 + else None + ) + # the loader accepts both (nframes, ngrid, 3) and the flattened + # (nframes, ngrid*3) layout; DeepDensity.eval only accepts the + # former, so normalize here and keep one contract at the eval call + grid = test_data["grid"].reshape(nframes, -1, 3) + + prediction = dp.eval( + coord, + box, + atype, + grid, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + mixed_type=mixed_type, + ).reshape(nframes, -1) + label = test_data["density"].reshape(nframes, -1) + + diff = prediction - label + errors: dict[str, tuple[float, float]] = { + "mae_density": (mae(diff), diff.size), + "rmse_density": (rmse(diff), diff.size), + } + + if context.detail_path is not None: + _write_per_frame_details( + context, + suffix="density", + reference=label, + prediction=prediction, + ) + + return errors diff --git a/deepmd/jax/atomic_model/dp_atomic_model.py b/deepmd/jax/atomic_model/dp_atomic_model.py index 8db31d6d5d..2ef1dc5eb9 100644 --- a/deepmd/jax/atomic_model/dp_atomic_model.py +++ b/deepmd/jax/atomic_model/dp_atomic_model.py @@ -40,6 +40,8 @@ class jax_atomic_model(dpmodel_atomic_model): """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls = None + """No long-range module is implemented in this backend.""" def forward_common_atomic( self, diff --git a/deepmd/jax/entrypoints/train.py b/deepmd/jax/entrypoints/train.py index 61d415a340..e86e5d4c8c 100644 --- a/deepmd/jax/entrypoints/train.py +++ b/deepmd/jax/entrypoints/train.py @@ -183,7 +183,7 @@ def factory( type_map = task_model.get_type_map() ipt_type_map = type_map if len(type_map) > 0 else None train_data = get_data( - dict(task_config.training_data_params), + task_config.training_data_params, task_model.get_rcut(), ipt_type_map, None, @@ -291,7 +291,7 @@ def update_sel( for task_config in iter_training_task_configs(jdata): type_map = task_config.model_params.get("type_map") train_data = get_data( - dict(task_config.training_data_params), + task_config.training_data_params, 0, # not used type_map, None, # not used diff --git a/deepmd/main.py b/deepmd/main.py index 1206b2fa0e..6d967553fc 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -369,6 +369,13 @@ def main_parser() -> argparse.ArgumentParser: "always freezes to 'graph', since the dense lower is deprecated in this " "backend -- requesting 'nlist' for one logs a warning and is overridden.", ) + parser_frz.add_argument( + "--compress", + action="store_true", + help="(Supported backend: PyTorch Exportable) Compress the model before " + "export using the default compression settings. Eligible .pt2 models " + "use the accelerated compressed graph without an intermediate export.", + ) # * test script ******************************************************************** parser_tst = subparsers.add_parser( @@ -1098,6 +1105,12 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: """ parser = main_parser() parsed_args = parser.parse_args(args=args) + if ( + parsed_args.command == "freeze" + and parsed_args.compress + and BACKEND_TABLE[parsed_args.backend] != "pytorch-exportable" + ): + parser.error("freeze --compress requires --pt-expt (PyTorch Exportable)") if parsed_args.command is None: parser.print_help() else: diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index cf15e99192..7495c1cc06 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -43,6 +43,7 @@ ) from deepmd.pd.utils.dataloader import ( DpLoaderSet, + create_dataset, ) from deepmd.pd.utils.env import ( DEVICE, @@ -66,9 +67,14 @@ from deepmd.utils.compat import ( update_deepmd_input, ) +from deepmd.utils.data_group import ( + GroupedDataset, +) +from deepmd.utils.data_path import ( + process_systems, +) from deepmd.utils.data_system import ( get_data, - process_systems, ) from deepmd.utils.model_preset import ( expand_model_preset, @@ -104,19 +110,11 @@ def prepare_trainer_input_single( data_dict_single: dict[str, Any], rank: int = 0, seed: int | None = None, - ) -> tuple[DpLoaderSet, DpLoaderSet | None, DPPath | None]: + ) -> tuple[ + DpLoaderSet | GroupedDataset, DpLoaderSet | GroupedDataset | None, DPPath | None + ]: training_dataset_params = data_dict_single["training_data"] validation_dataset_params = data_dict_single.get("validation_data") - validation_systems = ( - validation_dataset_params["systems"] if validation_dataset_params else None - ) - training_systems = training_dataset_params["systems"] - trn_patterns = training_dataset_params.get("rglob_patterns", None) - training_systems = process_systems(training_systems, patterns=trn_patterns) - if validation_systems is not None: - val_patterns = validation_dataset_params.get("rglob_patterns", None) - validation_systems = process_systems(validation_systems, val_patterns) - # stat files stat_file_path_single = data_dict_single.get("stat_file") if rank != 0: @@ -134,18 +132,16 @@ def prepare_trainer_input_single( # avoid the same batch sequence among devices rank_seed = [rank, seed % (2**32)] if seed is not None else None validation_data_single = ( - DpLoaderSet( - validation_systems, - validation_dataset_params["batch_size"], + create_dataset( + validation_dataset_params, model_params_single["type_map"], seed=rank_seed, ) - if validation_systems + if validation_dataset_params is not None else None ) - train_data_single = DpLoaderSet( - training_systems, - training_dataset_params["batch_size"], + train_data_single = create_dataset( + training_dataset_params, model_params_single["type_map"], seed=rank_seed, ) @@ -552,7 +548,10 @@ def change_bias( nbatches, ) updated_model = training.model_change_out_bias( - model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode + model_to_change, + sampled_data, + _bias_adjust_mode=bias_adjust_mode, + preset_out_bias={}, ) if not multi_task: diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index 8a8aaefd48..d3fdaed5fe 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -48,6 +48,7 @@ check_preset_out_bias, normalize_preset_out_bias, preset_out_bias_rows, + recover_nonfinite_preset_bias, remap_preset_out_bias, ) @@ -501,6 +502,7 @@ def change_out_bias( stat_file_path: DPPath | None = None, bias_adjust_mode: str = "change-by-statistic", observed_type: list[str] | None = None, + preset_out_bias: dict | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -524,6 +526,14 @@ def change_out_bias( ('set-by-statistic') or their stored bias ('change-by-statistic'); no statistics are computed, read or written for such an output, whose output std keeps its stored value. + preset_out_bias : dict, optional + An explicit downstream preset. Assigned outputs use this bias and + skip fitting; other outputs fit residuals against the stored bias. + The model preset is used when omitted, preserving training and + fine-tuning behavior. Pass an empty dict to fit from the data + without reapplying the model's stored preset. + Before residual fitting, non-finite stored cells covered by the + checkpoint preset are recovered; finite adjusted cells are retained. stat_file_path : Optional[DPPath] The path to the stat file. observed_type : list[str], optional @@ -533,17 +543,34 @@ def change_out_bias( if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"): raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) change = bias_adjust_mode == "change-by-statistic" + preset = ( + normalize_preset_out_bias(preset_out_bias, self.type_map) + if preset_out_bias is not None + else self.preset_out_bias + ) distinguish_types = self.get_compute_stats_distinguish_types() - check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types) + check_preset_out_bias(preset, self.bias_keys, distinguish_types) sampled = sample_merged + # A per-call override must not propagate a corrupt checkpoint baseline. + if change and preset_out_bias is not None and self.preset_out_bias: + recovered = recover_nonfinite_preset_bias( + self.preset_out_bias, + self.type_map, + to_numpy_array(self.out_bias), + self.bias_keys, + [self.atomic_output_def()[kk].size for kk in self.bias_keys], + ) + self._store_out_stat( + {kk: to_paddle_tensor(vv) for kk, vv in recovered.items()} + ) # === Step 1. Outputs fixed by the preset === fitted_keys = self.bias_keys - if self.preset_out_bias: + if preset: if observed_type is None: sampled = sample_merged() if callable(sample_merged) else sample_merged observed_type = collect_observed_types(sampled, self.type_map) rows = preset_out_bias_rows( - self.preset_out_bias, + preset, self.type_map, observed_type, to_numpy_array(self.out_bias), diff --git a/deepmd/pd/model/descriptor/dpa3.py b/deepmd/pd/model/descriptor/dpa3.py index 9a83fe56c4..77be53c6dd 100644 --- a/deepmd/pd/model/descriptor/dpa3.py +++ b/deepmd/pd/model/descriptor/dpa3.py @@ -215,6 +215,7 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.tebd_dim, precision=precision, seed=child_seed(seed, 3), + trainable=trainable, ) # 100 is a conservative upper bound self.spin_embedding = TypeEmbedNet( @@ -222,12 +223,14 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.tebd_dim, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) self.mix_cs_mlp = MLPLayer( 2 * self.tebd_dim, self.tebd_dim, precision=precision, seed=child_seed(seed, 5), + trainable=trainable, ) else: self.chg_embedding = None @@ -255,11 +258,14 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") ) - # set trainable - for param in self.parameters(): - param.requires_grad = trainable + self._apply_trainable() self.compress = False + def _apply_trainable(self) -> None: + """Apply the descriptor-level trainable setting to every parameter.""" + for param in self.parameters(): + param.stop_gradient = not self.trainable + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.rcut @@ -354,6 +360,11 @@ def share_params( assert self.__class__ == base_class.__class__, ( "Only descriptors of the same type can share params!" ) + if self.trainable != base_class.trainable: + raise ValueError( + "DPA3 descriptors must use the same trainable setting before " + "sharing parameters." + ) # For DPA3 descriptors, the user-defined share-level # shared_level: 0 # share all parameters in type_embedding, repflow @@ -384,6 +395,9 @@ def change_type_map( remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) self.type_map = type_map self.type_embedding.change_type_map(type_map=type_map) + # Type remapping replaces the first embedding matrix with a newly + # created parameter, so restore the descriptor-level frozen state. + self._apply_trainable() self.exclude_types = map_pair_exclude_types(self.exclude_types, remap_index) self.ntypes = len(type_map) repflow = self.repflows @@ -553,6 +567,9 @@ def t_cvt(xx: Any) -> paddle.Tensor: obj.repflows.layers = paddle.nn.LayerList( [RepFlowLayer.deserialize(layer) for layer in repflow_layers] ) + # Deserialization replaces several sublayers after construction, so apply + # the descriptor-level setting again to their newly registered parameters. + obj._apply_trainable() return obj def forward( diff --git a/deepmd/pd/model/model/__init__.py b/deepmd/pd/model/model/__init__.py index 59dce4f89b..94897bcfd1 100644 --- a/deepmd/pd/model/model/__init__.py +++ b/deepmd/pd/model/model/__init__.py @@ -64,6 +64,11 @@ def _get_standard_model_components( def get_standard_model(model_params: dict) -> BaseModel: + if model_params.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index f5beb501fc..3fa293cf64 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -231,6 +231,7 @@ def change_out_bias( self, merged: list[dict] | Callable[[], list[dict]], bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. @@ -248,10 +249,16 @@ def change_out_bias( 'change-by-statistic' : perform predictions on labels of target dataset, and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. + preset_out_bias : dict, optional + Per-call preset override. Omit to use the configured preset; + pass an empty dict to fit without fixing preset outputs. + Non-finite stored cells covered by the checkpoint preset are + recovered before residual fitting; finite values are retained. """ self.atomic_model.change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, ) def forward_common_lower( diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index a33908c632..f6464c43d0 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -24,15 +24,11 @@ from paddle.framework import ( core, ) -from paddle.io import ( - DataLoader, -) from deepmd.common import ( symlink_prefix_files, ) from deepmd.dpmodel.utils import ( - compute_total_numb_batch, resolve_model_prob, resolve_model_prob_from_epochs, ) @@ -59,7 +55,7 @@ ) from deepmd.pd.utils.dataloader import ( BufferedIterator, - get_sampler_from_params, + create_dataloader, ) from deepmd.pd.utils.env import ( CINN, @@ -67,7 +63,6 @@ DEFAULT_PRECISION, DEVICE, JIT, - NUM_WORKERS, SAMPLER_RECORD, enable_prim, ) @@ -185,23 +180,10 @@ def get_data_loader( def get_dataloader_and_buffer( _data: Any, _params: dict[str, Any] ) -> tuple[Any, Any]: - _sampler = get_sampler_from_params(_data, _params) - if _sampler is None: - log.warning( - "Sampler not specified!" - ) # None sampler will lead to a premature stop iteration. Replacement should be True in attribute of the sampler to produce expected number of items in one iteration. - _dataloader = DataLoader( - _data, - batch_sampler=paddle.io.BatchSampler( - sampler=_sampler, - drop_last=False, - ), - num_workers=NUM_WORKERS - if dist.is_available() - else 0, # setting to 0 diverges the behavior of its iterator; should be >=1 - collate_fn=lambda batch: batch[0], # prevent extra conversion + _dataloader = create_dataloader( + _data, _params, seed=_training_params.get("seed") ) - _data_buffered = BufferedIterator(iter(_dataloader)) + _data_buffered = BufferedIterator(_dataloader) return _dataloader, _data_buffered training_dataloader, training_data_buffered = get_dataloader_and_buffer( @@ -336,17 +318,9 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: self.validation_data, self.valid_numb_batch, ) = get_data_loader(training_data, validation_data, training_params) - training_data.print_summary( - "training", - to_numpy_array(self.training_dataloader.batch_sampler.sampler.weights), - ) + self.training_dataloader.print_summary("training") if validation_data is not None: - validation_data.print_summary( - "validation", - to_numpy_array( - self.validation_dataloader.batch_sampler.sampler.weights - ), - ) + self.validation_dataloader.print_summary("validation") else: ( self.training_dataloader, @@ -385,25 +359,15 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: training_params["data_dict"][model_key], ) - training_data[model_key].print_summary( - f"training in {model_key}", - to_numpy_array( - self.training_dataloader[ - model_key - ].batch_sampler.sampler.weights - ), + self.training_dataloader[model_key].print_summary( + f"training in {model_key}" ) if ( validation_data is not None and validation_data[model_key] is not None ): - validation_data[model_key].print_summary( - f"validation in {model_key}", - to_numpy_array( - self.validation_dataloader[ - model_key - ].batch_sampler.sampler.weights - ), + self.validation_dataloader[model_key].print_summary( + f"validation in {model_key}" ) per_task_total = [] @@ -415,13 +379,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: ) if self.num_epoch <= 0: raise ValueError("training.num_epoch must be positive.") - sampler_weights = to_numpy_array( - self.training_dataloader.batch_sampler.sampler.weights - ) - total_numb_batch = compute_total_numb_batch( - training_data.index, - sampler_weights, - ) + total_numb_batch = self.training_dataloader.epoch_length if total_numb_batch <= 0: raise ValueError( "Total number of training batches must be positive." @@ -441,16 +399,8 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: "are mutually exclusive." ) for model_key in self.model_keys: - sampler_weights = to_numpy_array( - self.training_dataloader[ - model_key - ].batch_sampler.sampler.weights - ) per_task_total.append( - compute_total_numb_batch( - training_data[model_key].index, - sampler_weights, - ) + self.training_dataloader[model_key].epoch_length ) ( self.model_prob, @@ -798,6 +748,25 @@ def single_model_finetune( self.profiling_file = training_params.get("profiling_file", "timeline.json") def run(self) -> None: + """Train while releasing native readers and producer threads on exit.""" + try: + self._run_training() + finally: + pipelines = ( + self.training_data, + self.validation_data, + self.training_dataloader, + self.validation_dataloader, + ) + for pipeline in pipelines: + sources = ( + pipeline.values() if isinstance(pipeline, dict) else (pipeline,) + ) + for source in sources: + if source is not None: + source.close() + + def _run_training(self) -> None: fout = ( open( self.disp_file, @@ -1184,9 +1153,7 @@ def get_data( batch_data = next(iter(self.training_data)) except StopIteration: # Refresh the status of the dataloader to start from a new epoch - self.training_data = BufferedIterator( - iter(self.training_dataloader) - ) + self.training_data = BufferedIterator(self.training_dataloader) batch_data = next(iter(self.training_data)) else: if self.validation_data is None: @@ -1194,9 +1161,7 @@ def get_data( try: batch_data = next(iter(self.validation_data)) except StopIteration: - self.validation_data = BufferedIterator( - iter(self.validation_dataloader) - ) + self.validation_data = BufferedIterator(self.validation_dataloader) batch_data = next(iter(self.validation_data)) else: if is_train: @@ -1205,7 +1170,7 @@ def get_data( except StopIteration: # Refresh the status of the dataloader to start from a new epoch self.training_data[task_key] = BufferedIterator( - iter(self.training_dataloader[task_key]) + self.training_dataloader[task_key] ) batch_data = next(iter(self.training_data[task_key])) else: @@ -1215,7 +1180,7 @@ def get_data( batch_data = next(iter(self.validation_data[task_key])) except StopIteration: self.validation_data[task_key] = BufferedIterator( - iter(self.validation_dataloader[task_key]) + self.validation_dataloader[task_key] ) batch_data = next(iter(self.validation_data[task_key])) @@ -1455,11 +1420,16 @@ def model_change_out_bias( _model: Any, _sample_func: Any, _bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: old_bias = deepcopy(_model.get_out_bias()) + preset_kwargs = ( + {"preset_out_bias": preset_out_bias} if preset_out_bias is not None else {} + ) _model.change_out_bias( _sample_func, bias_adjust_mode=_bias_adjust_mode, + **preset_kwargs, ) new_bias = deepcopy(_model.get_out_bias()) diff --git a/deepmd/pd/utils/dataloader.py b/deepmd/pd/utils/dataloader.py index 090d1713f5..d0e78ce4d8 100644 --- a/deepmd/pd/utils/dataloader.py +++ b/deepmd/pd/utils/dataloader.py @@ -7,10 +7,14 @@ Iterable, Iterator, ) +from functools import ( + cached_property, +) from multiprocessing.dummy import ( Pool, ) from threading import ( + Event, Thread, ) from typing import ( @@ -42,15 +46,29 @@ ) from deepmd.pd.utils.utils import ( mix_entropy, + to_numpy_array, ) from deepmd.utils import random as dp_random from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_config import ( + DataConfig, + data_group_configs, + resolve_data_groups, +) +from deepmd.utils.data_group import ( + GroupedDataset, + GroupedTrainingDataLoader, + TrainingDataLoader, +) +from deepmd.utils.data_sampling import ( + data_child_seeds, + global_batch_counts, + system_probabilities, +) from deepmd.utils.data_system import ( print_summary, - prob_sys_size_ext, - process_sys_probs, ) log = logging.getLogger(__name__) @@ -228,6 +246,22 @@ def __next__(self) -> Any: def __len__(self) -> int: return len(self.dataloaders) + @cached_property + def mean_batch_nframes(self) -> np.ndarray: + """Count frames per native batch, including tails and distributed padding.""" + return np.asarray( + [ + ( + sampler.num_samples + if isinstance(sampler, DistributedBatchSampler) + else len(sampler.sampler) + ) + / count + for sampler, count in zip(self.sampler_list, self.index, strict=True) + ], + dtype=np.float64, + ) + def __getitem__(self, idx: int) -> dict: # log.warning(str(paddle.distributed.get_rank())+" idx: "+str(idx)+" index: "+str(self.index[idx])) try: @@ -277,19 +311,38 @@ def __init__( self._queue = queue self._source = source # Main DL iterator self._max_len = max_len + self._stop_requested = Event() def run(self) -> None: for item in self._source: - self._queue.put(item) # Blocking if the queue is full + if self._stop_requested.is_set(): + return + if not self._put(item): + return # Signal the consumer we are done. - self._queue.put(_sentinel) + self._put(_sentinel) + + def _put(self, item: Any) -> bool: + """Wait for buffer capacity while allowing the owner to stop loading.""" + while not self._stop_requested.is_set(): + try: + self._queue.put(item, timeout=0.1) + return True + except queue.Full: + continue + return False + + def close(self) -> None: + """Stop producing batches and wait for the active read to complete.""" + self._stop_requested.set() + self.join() class BufferedIterator: def __init__(self, iterable: Iterable) -> None: self._queue = queue.Queue(QUEUESIZE) - self._iterable = iterable + self._iterable = iter(iterable) self._consumer = None self.start_time = time.time() @@ -308,6 +361,8 @@ def __len__(self) -> int: return self.total def __next__(self) -> Any: + if self._iterable is None: + raise StopIteration # Create consumer if not created yet if self._consumer is None: self._create_consumer() @@ -330,9 +385,19 @@ def __next__(self) -> Any: if isinstance(item, Exception): raise item if item is _sentinel: + self.close() raise StopIteration return item + def close(self) -> None: + """Release buffered batches before the owner closes the native loader.""" + if self._consumer is not None: + self._consumer.close() + self._consumer = None + self._iterable = None + while not self._queue.empty(): + self._queue.get_nowait() + def collate_batch(batch: list[dict]) -> dict: example = batch[0] @@ -352,21 +417,39 @@ def collate_batch(batch: list[dict]) -> dict: return result +class DpDataLoader(DataLoader): + """Own native reader threads and persistent per-system worker cursors.""" + + def __iter__(self) -> Iterator[dict[str, Any]]: + self._active_iterator = super().__iter__() + return self._active_iterator + + def close(self) -> None: + """Release native threads or workers, including a partially consumed pass.""" + iterator = getattr(self, "_active_iterator", None) + if iterator is not None: + iterator._try_shutdown_all() + self._active_iterator = None + self._iterator = None + + def __del__(self) -> None: + """Release workers if construction or initialization ends before training.""" + self.close() + + def get_weighted_sampler( - training_data: DpLoaderSet, prob_style: str, sys_prob: bool = False + training_data: DpLoaderSet, prob_style: str | list[float], sys_prob: bool = False ) -> WeightedRandomSampler: - if sys_prob is False: - if prob_style == "prob_uniform": - prob_v = 1.0 / float(training_data.__len__()) - probs = [prob_v for ii in range(training_data.__len__())] - else: # prob_sys_size;A:B:p1;C:D:p2 or prob_sys_size = prob_sys_size;0:nsys:1.0 - if prob_style == "prob_sys_size": - style = f"prob_sys_size;0:{len(training_data)}:1.0" - else: - style = prob_style - probs = prob_sys_size_ext(style, len(training_data), training_data.index) - else: - probs = process_sys_probs(prob_style, training_data.index) + probs = system_probabilities( + training_data.index, + sys_probs=prob_style if sys_prob else None, + auto_prob="prob_sys_size" if sys_prob else prob_style, + batch_nframes=( + training_data.mean_batch_nframes + if not sys_prob and prob_style == "prob_uniform" + else None + ), + ) log.debug("Generated weighted sampler with prob array: " + str(probs)) # training_data.total_batch is the size of one epoch, you can increase it to avoid too many rebuilding of iteraters len_sampler = training_data.total_batch * max(env.NUM_WORKERS, 1) @@ -390,3 +473,65 @@ def get_sampler_from_params( else: _sampler = get_weighted_sampler(_data, "prob_sys_size") return _sampler + + +def create_dataset( + config: DataConfig, type_map: list[str], *, seed: int | list[int] | None = None +) -> DpLoaderSet | GroupedDataset: + """Construct NumPy loader groups through the shared path resolver.""" + groups = resolve_data_groups(config, formats=("deepmd",)) + datasets = [ + DpLoaderSet(list(group.systems), group.batch_size, type_map, seed=seed) + for group in groups + ] + return ( + datasets[0] + if len(datasets) == 1 + else GroupedDataset( + datasets, data_group_configs(config), [group.prob for group in groups] + ) + ) + + +def create_dataloader( + data: DpLoaderSet | GroupedDataset, + config: DataConfig, + *, + seed: int | None = None, +) -> TrainingDataLoader | GroupedTrainingDataLoader: + """Bind native Paddle loaders to the shared batch-composition interface.""" + if isinstance(data, GroupedDataset): + seeds = data_child_seeds( + 0 if seed is None and dist.get_world_size() > 1 else seed, + len(data.children) + 1, + ) + children = [ + create_dataloader(child, params, seed=child_seed) + for child, params, child_seed in zip( + data.children, data.configs, seeds[1:], strict=True + ) + ] + return GroupedTrainingDataLoader(data, children, seeds[0]) + params = config[0] if isinstance(config, list) else config + sampler = get_sampler_from_params(data, params) + num_workers = env.NUM_WORKERS if dist.is_available() else 0 + loader = DpDataLoader( + data, + batch_sampler=BatchSampler(sampler=sampler, drop_last=False), + num_workers=num_workers, + persistent_workers=( + num_workers > 0 + and params.get("auto_prob") == "prob_uniform" + and params.get("sys_probs") is None + ), + collate_fn=lambda batch: batch[0], + ) + return TrainingDataLoader( + data, + loader, + to_numpy_array(sampler.weights), + lambda: data.index, + sampling_sizes=lambda: global_batch_counts( + [len(system) for system in data.systems], data.batch_sizes + ), + ) diff --git a/deepmd/pd/utils/dataset.py b/deepmd/pd/utils/dataset.py index bd8a0bc91e..90a87edcf6 100644 --- a/deepmd/pd/utils/dataset.py +++ b/deepmd/pd/utils/dataset.py @@ -57,4 +57,5 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N default=data_item["default"], dtype=data_item["dtype"], output_natoms_for_type_sel=data_item["output_natoms_for_type_sel"], + special_shape=getattr(data_item, "special_shape", None), ) diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index b3cd33798a..c6860678b0 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -125,8 +125,7 @@ def _strip_shape_assertions(graph_module: torch.nn.Module) -> None: SeZM lower inputs intentionally keep extended-atom and local-atom axes independent: regular exports pass ghost coordinates through ``coord`` while - ``atype`` remains local-only, and spin exports slice both ``nall`` and - ``nloc`` after virtual atom expansion. ``torch.export`` may turn these valid + ``atype`` remains local-only. ``torch.export`` may turn these valid dynamic cases into deferred ``Ne(nall, nloc)`` assertions. """ graph = graph_module.graph @@ -493,11 +492,8 @@ def _make_sample_inputs( ) -> tuple[torch.Tensor | None, ...]: """Build representative ``forward_common_lower`` inputs for tracing. - Three lower ABIs are produced, selected by ``model.export_lower_input_kind()`` - and whether the model carries spin: + The edge-schema inputs depend on whether the model carries spin: - - virtual spin (``nlist``): the DeepSpin extended-input signature, since the - graph expands virtual atoms internally; - native spin (``edge_vec``): the energy edge schema plus the owned-atom spins (the first ``nloc`` extended rows, where ``mapping`` is identity); - energy (``edge_vec``): the plain single-domain edge schema. @@ -512,17 +508,6 @@ def _make_sample_inputs( aparam, charge_spin, ) = _build_sample_extended(model, nframes, nloc, device, has_spin) - if has_spin and model.export_lower_input_kind() == "nlist": - return ( - ext_coord, - ext_atype, - ext_spin, - nlist_t, - mapping_t, - fparam, - aparam, - charge_spin, - ) formatted_nlist: torch.Tensor = model.format_nlist(ext_coord, ext_atype, nlist_t) edge_schema = edge_schema_from_extended( ext_coord, @@ -531,6 +516,7 @@ def _make_sample_inputs( mapping_t, ) if has_spin: + # Local spins have a frame stride independent of ghost storage. return ( edge_schema.coord, edge_schema.atype, @@ -538,7 +524,7 @@ def _make_sample_inputs( edge_schema.edge_vec, edge_schema.edge_scatter_index, edge_schema.edge_mask, - ext_spin[:, :nloc], + ext_spin[:, :nloc].contiguous(), fparam, aparam, charge_spin, @@ -681,38 +667,12 @@ def _build_dynamic_shapes( ) -> tuple: """Build positional dynamic-shape constraints for the traced lower input. - The lower ABI is recovered from the sample structure: a floating-point - tensor at index 2 is the extended spin of the deepspin-scheme nlist contract, - while an integer ``edge_index`` there marks the edge contract. A native-spin - edge sample carries the extra per-local-atom spin tensor, giving it ten - positional entries against the energy contract's nine. + A native-spin edge sample carries the extra per-local-atom spin tensor, + giving it ten positional entries against the energy contract's nine. """ nframes_dim = torch.export.Dim("nframes", min=1) nloc_dim = torch.export.Dim("nloc", min=1) nedge_dim = torch.export.Dim("nedge", min=2) - is_nlist_spin = ( - len(sample_inputs) >= 3 - and sample_inputs[2] is not None - and sample_inputs[2].is_floating_point() - ) - if is_nlist_spin: - nall_dim = torch.export.Dim("nall", min=4) - fparam = sample_inputs[5] - aparam = sample_inputs[6] - charge_spin = sample_inputs[7] if len(sample_inputs) == 8 else None - shapes = ( - {0: nframes_dim, 1: nall_dim}, # extended_coord - {0: nframes_dim, 1: nall_dim}, # extended_atype - {0: nframes_dim, 1: nall_dim}, # extended_spin - {0: nframes_dim, 1: nloc_dim}, # nlist - {0: nframes_dim, 1: nall_dim}, # mapping - {0: nframes_dim} if fparam is not None else None, - {0: nframes_dim, 1: nloc_dim} if aparam is not None else None, - ) - if len(sample_inputs) == 8: - shapes = (*shapes, {0: nframes_dim} if charge_spin is not None else None) - return shapes - nall_dim = torch.export.Dim("nall", min=1) edge_shapes = ( {0: nframes_dim, 1: nall_dim}, # extended_coord: (nframes, nall, 3) @@ -917,6 +877,7 @@ def freeze_sezm_to_pt2( device: torch.device | None = None, head: str | None = None, atomic_virial: bool = True, + with_comm: bool = True, ) -> None: """Freeze a SeZM checkpoint into an AOTInductor ``.pt2`` archive. @@ -938,6 +899,9 @@ def freeze_sezm_to_pt2( default: the edge-force scatter assembles the per-atom virial as a free by-product of the single backward, so exporting it carries no compute cost. + with_comm + Whether to include a communication artifact when the model supports + cross-rank message passing. Disable for single-process inference. Notes ----- @@ -955,6 +919,7 @@ def freeze_sezm_to_pt2( target_device=target_device, head=head, atomic_virial=atomic_virial, + with_comm=with_comm, ) @@ -965,6 +930,7 @@ def _freeze_sezm_to_pt2( target_device: torch.device, head: str | None, atomic_virial: bool, + with_comm: bool, ) -> None: """Build one AOTInductor archive under an established kernel policy.""" from torch._inductor import ( @@ -1036,8 +1002,7 @@ def _freeze_sezm_to_pt2( ) # Each model's exportable signature matches its sample tuple positionally - # (energy / native-spin edge ABI, or virtual-spin nlist ABI), so a single - # splat covers all three contracts. + # (energy or native-spin edge ABI). log.info("Tracing the lower graph on CPU (make_fx)...") traced = model.forward_common_lower_exportable(*sample_inputs_cpu) @@ -1080,17 +1045,17 @@ def _freeze_sezm_to_pt2( # Second artifact: the LAMMPS multi-rank with-comm graph. It threads the # eight border_op communication tensors so cross-rank ghost features are - # exchanged between interaction blocks. Gated on the edge_vec lower contract - # (energy and native spin), so virtual spin (nlist interface) is excluded - # and falls back to single-rank. Bridging models DO take this path: they + # exchanged between interaction blocks. Bridging models take this path: they # report supports_edge_parallel()=True since the Source Freeze Propagation # gate's per-node partials are completed across ranks by # DescrptSeZM._gate_partial_exchange (issue #5906). - with_comm = ( - model.export_lower_input_kind() == "edge_vec" and model.supports_edge_parallel() + has_comm_artifact = ( + with_comm + and model.export_lower_input_kind() == "edge_vec" + and model.supports_edge_parallel() ) with_comm_bytes: bytes | None = None - if with_comm: + if has_comm_artifact: log.info( "Compiling the parallel with-comm artifact (second AOTInductor " "compilation)..." @@ -1106,7 +1071,7 @@ def _freeze_sezm_to_pt2( output_keys=output_keys, is_spin=is_spin, do_atomic_virial=atomic_virial, - has_comm_artifact=with_comm, + has_comm_artifact=has_comm_artifact, ) with zipfile.ZipFile(out_path_str, "a") as zf: zf.writestr("model/extra/metadata.json", json.dumps(metadata)) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 9d54519e9d..43ba0bfa61 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -61,6 +61,7 @@ ) from deepmd.pt.utils.dataloader import ( DpLoaderSet, + create_dataset, ) from deepmd.pt.utils.env import ( DEVICE, @@ -70,7 +71,6 @@ ) from deepmd.pt.utils.lmdb_dataset import ( LmdbDataset, - is_lmdb, ) from deepmd.pt.utils.multi_task import ( preprocess_shared_params, @@ -90,10 +90,15 @@ from deepmd.utils.data import ( min_pair_dist_requirement, ) -from deepmd.utils.data_system import ( - get_data, +from deepmd.utils.data_group import ( + GroupedDataset, +) +from deepmd.utils.data_path import ( process_systems, ) +from deepmd.utils.data_system import ( + get_neighbor_stat_data, +) from deepmd.utils.model_preset import ( expand_model_preset, ) @@ -155,8 +160,8 @@ def prepare_trainer_input_single( rank: int = 0, seed: int | None = None, ) -> tuple[ - DpLoaderSet | LmdbDataset, - DpLoaderSet | LmdbDataset | None, + DpLoaderSet | LmdbDataset | GroupedDataset, + DpLoaderSet | LmdbDataset | GroupedDataset | None, StatFileSpec, ]: # get data modifier @@ -167,11 +172,6 @@ def prepare_trainer_input_single( training_dataset_params = data_dict_single["training_data"] validation_dataset_params = data_dict_single.get("validation_data") - validation_systems = ( - validation_dataset_params["systems"] if validation_dataset_params else None - ) - training_systems = training_dataset_params["systems"] - stat_file_spec = StatFileSpec( data_dict_single.get("stat_file"), data_dict_single.get("stat_file_mode", "update"), @@ -179,56 +179,22 @@ def prepare_trainer_input_single( rank_seed = [rank, seed % (2**32)] if seed is not None else None - def _make_dp_loader_set( - systems: str | list[str], - dataset_params: dict[str, Any], - ) -> DpLoaderSet: - """Create a DpLoaderSet from systems with pattern expansion.""" - patterns = dataset_params.get("rglob_patterns") - systems = process_systems(systems, patterns=patterns) - return DpLoaderSet( - systems, - dataset_params["batch_size"], + train_data_single = create_dataset( + training_dataset_params, + model_params_single["type_map"], + seed=rank_seed, + modifier=modifier, + ) + validation_data_single = ( + create_dataset( + validation_dataset_params, model_params_single["type_map"], seed=rank_seed, modifier=modifier, ) - - # LMDB path: single string → LmdbDataset - if isinstance(training_systems, str) and is_lmdb(training_systems): - auto_prob = training_dataset_params.get("auto_prob", None) - train_data_single = LmdbDataset( - training_systems, - model_params_single["type_map"], - training_dataset_params["batch_size"], - auto_prob_style=auto_prob, - ) - if ( - validation_systems is not None - and isinstance(validation_systems, str) - and is_lmdb(validation_systems) - ): - validation_data_single = LmdbDataset( - validation_systems, - model_params_single["type_map"], - validation_dataset_params["batch_size"], - ) - elif validation_systems is not None: - validation_data_single = _make_dp_loader_set( - validation_systems, validation_dataset_params - ) - else: - validation_data_single = None - else: - # Standard npy path - train_data_single = _make_dp_loader_set( - training_systems, training_dataset_params - ) - validation_data_single = ( - _make_dp_loader_set(validation_systems, validation_dataset_params) - if validation_systems - else None - ) + if validation_dataset_params is not None + else None + ) return ( train_data_single, validation_data_single, @@ -399,21 +365,9 @@ def train( if not multi_task: type_map = config["model"].get("type_map") - training_systems = config["training"]["training_data"].get("systems") - if ( - training_systems is not None - and isinstance(training_systems, str) - and is_lmdb(training_systems) - ): - from deepmd.dpmodel.utils.lmdb_data import ( - make_neighbor_stat_data, - ) - - train_data = make_neighbor_stat_data(training_systems, type_map) - else: - train_data = get_data( - config["training"]["training_data"], 0, type_map, None - ) + train_data = get_neighbor_stat_data( + config["training"]["training_data"], type_map + ) config["model"], min_nbor_dist = BaseModel.update_sel( train_data, type_map, config["model"] ) @@ -421,26 +375,10 @@ def train( min_nbor_dist = {} for model_item in config["model"]["model_dict"]: type_map = config["model"]["model_dict"][model_item].get("type_map") - training_systems = config["training"]["data_dict"][model_item][ - "training_data" - ].get("systems") - if ( - training_systems is not None - and isinstance(training_systems, str) - and is_lmdb(training_systems) - ): - from deepmd.dpmodel.utils.lmdb_data import ( - make_neighbor_stat_data, - ) - - train_data = make_neighbor_stat_data(training_systems, type_map) - else: - train_data = get_data( - config["training"]["data_dict"][model_item]["training_data"], - 0, - type_map, - None, - ) + train_data = get_neighbor_stat_data( + config["training"]["data_dict"][model_item]["training_data"], + type_map, + ) config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = ( BaseModel.update_sel( train_data, type_map, config["model"]["model_dict"][model_item] @@ -637,7 +575,10 @@ def change_bias( min_pair_dist=0.0 if pair_filter is None else float(pair_filter.default), ) updated_model = training.model_change_out_bias( - model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode + model_to_change, + sampled_data, + _bias_adjust_mode=bias_adjust_mode, + preset_out_bias={}, ) if not multi_task: diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index f468379565..272e39227d 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -20,6 +20,9 @@ OutputVariableCategory, OutputVariableDef, ) +from deepmd.infer.deep_density import ( + DeepDensity, +) from deepmd.infer.deep_dipole import ( DeepDipole, ) @@ -69,17 +72,12 @@ GLOBAL_PT_FLOAT_PRECISION, RESERVED_PRECISION_DICT, ) -from deepmd.pt.utils.nv_nlist import ( - NvNeighborList, - is_nv_available, -) from deepmd.pt.utils.utils import ( to_numpy_array, to_torch_tensor, ) -from deepmd.pt_expt.utils.vesin_neighbor_list import ( - VesinNeighborList, - is_vesin_torch_available, +from deepmd.pt_expt.utils.auto_neighbor_list import ( + make_neighbor_list, ) from deepmd.utils.batch_size import ( RetrySignal, @@ -113,7 +111,7 @@ def _is_sezm_model_params(model_params: dict[str, Any]) -> bool: """Return whether the params describe a SeZM / DPA4 model.""" model_type = str(model_params.get("type", "")).lower() - if model_type in {"sezm", "dpa4", "sezm_spin"}: + if model_type in {"sezm", "dpa4"}: return True # canonical bridged spelling: linear_ener over [dpa4, inner_potential], # realized by the pt backend as a SeZMModel @@ -143,6 +141,15 @@ class DeepEval(DeepEvalBackend): neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional The ASE neighbor list class to produce the neighbor list. If None, the neighbor list will be built natively in the model. + nlist_backend : str, default: "auto" + Extended-coordinate neighbor search: ``"auto"``, ``"native"``, + ``"vesin"``, or ``"nv"``. Auto prefers Vesin for a single frame and + Toolkit-Ops for CUDA batches when available. + neighbor_graph_method : str, default: "auto" + SeZM/DPA4 inference search: ``"auto"``, ``"vesin"``, or ``"nv"``. + The strategy runs inside the model's native forward geometry stage. + neighbor_skin : float, default: 0.0 + Vesin candidate-search skin in angstrom. Zero disables Verlet caching. **kwargs : dict Keyword arguments. """ @@ -157,11 +164,15 @@ def __init__( head: str | int | None = None, no_jit: bool = False, nlist_backend: str = "auto", + neighbor_graph_method: str = "auto", + neighbor_skin: float = 0.0, **kwargs: Any, ) -> None: self.output_def = output_def self.model_path = model_file self.neighbor_list = neighbor_list + self._neighbor_graph_method = neighbor_graph_method + self._neighbor_skin = neighbor_skin # data modifier, populated only for frozen .pth models that carry one self.modifier = None if str(self.model_path).endswith(".pt"): @@ -279,25 +290,47 @@ def __init__( def _setup_nlist_backend(self, nlist_backend: str) -> None: """Resolve the neighbor-list construction strategy from a user choice. - ``"native"`` uses the dense all-pairs builder; ``"vesin"`` / ``"nv"`` - force the O(N) ``vesin.torch`` / ``nvalchemiops`` cell list (raising if - unavailable or the model/inputs are unsupported); ``"auto"`` picks the - first available O(N) builder (vesin, then nv) and otherwise falls back to - the native builder. Results are unchanged either way -- only the - neighbor-search cost differs. + ``"native"`` retains the model's unbuffered construction policy and + ignores ``neighbor_skin``. For self-built SeZM/DPA4 models this is a + batch-aware Vesin/NV strategy, with dense carry-all edges only when both + optional libraries are unavailable. Other models retain their native + extended-coordinate construction. ``"vesin"`` / ``"nv"`` force the + corresponding search (raising for unavailable or unsupported inputs); + ``"auto"`` selects by frame count and available libraries. Results are + unchanged either way -- only the neighbor-search cost differs. """ inner = self.dp.model["Default"] self_built = getattr(inner, "use_self_built_nlist", None) if callable(self_built) and self_built(): - # The model builds its own neighbor list and runs the native path; - # an external strategy would bypass it, so always use native. - log.debug( - "Ignoring nlist_backend=%r: %s uses its own built-in neighbor list.", - nlist_backend, - type(inner).__name__, - ) + # The strategy belongs to the model's geometry stage. Keeping the + # native forward entry preserves its compiled core computation. + configure = getattr(inner, "set_inference_neighbor_list", None) self._nlist_builder = None + if not callable(configure): + return + method = self._neighbor_graph_method + if method not in ("auto", "vesin", "nv"): + raise ValueError( + "PT SeZM/DPA4 supports neighbor_graph_method 'auto', " + f"'vesin', or 'nv'; got {method!r}." + ) + if method != "auto" and nlist_backend != "auto": + raise ValueError( + "Select only one of neighbor_graph_method and nlist_backend." + ) + if method == "auto": + method = nlist_backend + configure( + make_neighbor_list( + method, DEVICE, self._neighbor_skin, return_mode="edges" + ) + ) return + if self._neighbor_graph_method != "auto": + raise ValueError( + "neighbor_graph_method applies to PT SeZM/DPA4 models; " + "use nlist_backend for this model." + ) if nlist_backend not in ("auto", "vesin", "nv", "native"): raise ValueError( f"Unknown nlist_backend '{nlist_backend}'; " @@ -325,7 +358,6 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None: unsupported = "non-energy models" ase_provided = self.neighbor_list is not None - builder = None if nlist_backend in ("vesin", "nv"): if unsupported is not None: raise ValueError( @@ -337,36 +369,11 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None: f"nlist_backend='{nlist_backend}' conflicts with an " "explicitly supplied ASE neighbor_list; pass only one." ) - if nlist_backend == "vesin": - if not is_vesin_torch_available(): - raise ImportError( - "nlist_backend='vesin' was requested but 'vesin.torch' " - "is not installed. Install it (`pip install " - "vesin[torch]`) or use nlist_backend='native' (or 'auto')." - ) - builder = VesinNeighborList() - elif DEVICE.type != "cuda": - raise ValueError( - "nlist_backend='nv' requires CUDA inference tensors; " - f"current DEVICE is {DEVICE!s}. Use nlist_backend='native' " - "(or 'auto') for CPU inference." - ) - elif not is_nv_available(): - raise ImportError( - "nlist_backend='nv' was requested but 'nvalchemi-toolkit-ops'" - " is not installed. Install it (`pip install " - "nvalchemi-toolkit-ops`) or use nlist_backend='native' " - "(or 'auto')." - ) - else: - builder = NvNeighborList() - elif nlist_backend == "auto" and unsupported is None and not ase_provided: - # Pick the first available O(N) builder; nv is GPU-only. - if is_vesin_torch_available(): - builder = VesinNeighborList() - elif is_nv_available() and DEVICE.type == "cuda": - builder = NvNeighborList() - self._nlist_builder = builder + self._nlist_builder = ( + make_neighbor_list(nlist_backend, DEVICE, self._neighbor_skin) + if unsupported is None and not ase_provided + else None + ) def get_rcut(self) -> float: """Get the cutoff radius of this model.""" @@ -440,11 +447,26 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepWFC elif "population" in model_output_type: return DeepPopulation + elif "density" in model_output_type and self._model_has_grid( + self.dp.model["Default"] + ): + # key on the grid capability rather than the output name, so a + # property fitting with property_name "density" still dispatches + # to DeepProperty below + return DeepDensity elif self.get_var_name() in model_output_type: return DeepProperty else: raise RuntimeError("Unknown model type") + @staticmethod + def _model_has_grid(model: Any) -> bool: + has_grid = getattr(model, "has_grid", None) + try: + return bool(has_grid()) if callable(has_grid) else False + except AttributeError: + return False + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. @@ -551,6 +573,41 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + if "grid" in kwargs and kwargs["grid"] is not None: + grid_input = np.array(kwargs["grid"]) + # the directional neighbor list is dense in ngrid x nall, so the + # batching size proxy must account for the grid extent, not just + # the atom count + ngrid = grid_input.size // (numb_test * 3) + out = self._eval_func( + self._eval_model_density, numb_test, max(natoms, ngrid) + )( + coords, + cells, + atom_types, + grid_input, + fparam, + aparam, + self._get_request_defs(atomic), + charge_spin, + ) + # _eval_model_density returns a 1-element tuple; execute_all unwraps + # it when auto batching is enabled, but with auto_batch_size=False + # the inner function is called directly and the tuple survives. + # request_defs is ignored on this path: the density output is a + # single fixed variable, so atomic=True returns the same dict as + # atomic=False. Kept for signature parity with the other + # _eval_model_* paths. + if isinstance(out, tuple): + (out,) = out + return {"density": out} + if "density" in self.output_def.var_defs and self._model_has_grid( + self.dp.model["Default"] + ): + raise ValueError( + "grid is required to evaluate a density model; " + "pass grid=... with shape (nframes, ngrid, 3)" + ) request_defs = self._get_request_defs(atomic) if "spin" not in kwargs or kwargs["spin"] is None: out = self._eval_func(self._eval_model, numb_test, natoms)( @@ -916,6 +973,86 @@ def _eval_model_spin( ) # this is kinda hacky return tuple(results) + def _eval_model_density( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + grid: np.ndarray, + fparam: np.ndarray | None, + aparam: np.ndarray | None, + request_defs: list[OutputVariableDef], + charge_spin: np.ndarray | None, + ) -> tuple[np.ndarray, ...]: + model = self.dp.to(DEVICE) + + nframes = coords.shape[0] + if len(atom_types.shape) == 1: + natoms = len(atom_types) + atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) + else: + natoms = len(atom_types[0]) + + coord_input = torch.tensor( + coords.reshape([nframes, natoms, 3]), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=DEVICE, + ) + type_input = torch.tensor(atom_types, dtype=torch.long, device=DEVICE) + grid_input = torch.tensor( + grid.reshape([nframes, -1, 3]), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=DEVICE, + ) + ngrid = grid_input.shape[1] + if cells is not None: + box_input = torch.tensor( + cells.reshape([nframes, 3, 3]), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=DEVICE, + ) + else: + box_input = None + if fparam is not None: + fparam_input = to_torch_tensor( + fparam.reshape(nframes, self.get_dim_fparam()) + ) + else: + fparam_input = None + if aparam is not None: + aparam_input = to_torch_tensor( + aparam.reshape(nframes, natoms, self.get_dim_aparam()) + ) + else: + aparam_input = None + if charge_spin is not None: + charge_spin_input = to_torch_tensor(charge_spin.reshape(nframes, 2)) + else: + charge_spin_input = None + + do_atomic_virial = any( + x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs + ) + batch_output = model( + coord_input, + type_input, + grid=grid_input, + box=box_input, + do_atomic_virial=do_atomic_virial, + fparam=fparam_input, + aparam=aparam_input, + charge_spin=charge_spin_input, + ) + if isinstance(batch_output, tuple): + batch_output = batch_output[0] + + results = [] + pt_name = "density" + density_shape = [nframes, ngrid] + out = batch_output[pt_name].reshape(density_shape).detach().cpu().numpy() + results.append(out) + return tuple(results) + def _get_output_shape( self, odef: OutputVariableDef, nframes: int, natoms: int ) -> list[int]: diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py index 7a47d13ae2..544320faf5 100644 --- a/deepmd/pt/loss/__init__.py +++ b/deepmd/pt/loss/__init__.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from .charge import ( + GridDensityLoss, +) from .dens import ( DeNSLoss, ) @@ -31,6 +34,7 @@ "EnergyHessianStdLoss", "EnergySpinLoss", "EnergyStdLoss", + "GridDensityLoss", "PopulationLoss", "PropertyLoss", "TaskLoss", diff --git a/deepmd/pt/loss/charge.py b/deepmd/pt/loss/charge.py new file mode 100644 index 0000000000..53b2250434 --- /dev/null +++ b/deepmd/pt/loss/charge.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch + +from deepmd.pt.loss.loss import ( + TaskLoss, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + GLOBAL_PT_FLOAT_PRECISION, +) +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +class GridDensityLoss(TaskLoss): + def __init__( + self, + starter_learning_rate: float = 1.0, + start_pref_d: float = 0.0, + limit_pref_d: float = 0.0, + inference: bool = False, + **kwargs: Any, + ) -> None: + r"""Construct a layer to compute loss on grid density. + + The residual is masked by the grid-point mask (excluded points do + not contribute) and reduced per frame before averaging over the + batch, consistent with the other losses in the package. + + Parameters + ---------- + starter_learning_rate : float + The learning rate at the start of the training. + start_pref_d : float + The prefactor of charge density loss at the start of the training. + limit_pref_d : float + The prefactor of charge density loss at the end of the training. + inference : bool + If true, it will output all losses found in output, ignoring the pre-factors. + **kwargs + Other keyword arguments. + """ + super().__init__() + self.starter_learning_rate = starter_learning_rate + self.has_d = (start_pref_d != 0.0 or limit_pref_d != 0.0) or inference + + self.start_pref_d = start_pref_d + self.limit_pref_d = limit_pref_d + self.inference = inference + + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float, + mae: bool = False, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: + """Return loss on grid density. + + Parameters + ---------- + input_dict : dict[str, torch.Tensor] + Model inputs. + model : torch.nn.Module + Model to be used to output the predictions. + label : dict[str, torch.Tensor] + Labels. + natoms : int + The local atom number. + + Returns + ------- + model_pred: dict[str, torch.Tensor] + Model predictions. + loss: torch.Tensor + Loss for model to minimize. + more_loss: dict[str, torch.Tensor] + Other losses for display. + """ + model_pred = model(**input_dict) + coef = learning_rate / self.starter_learning_rate + pref_d = self.limit_pref_d + (self.start_pref_d - self.limit_pref_d) * coef + + loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] + more_loss = {} + # more_loss['log_keys'] = [] # showed when validation on the fly + # more_loss['test_keys'] = [] # showed when doing dp test + if self.has_d and "density" in model_pred and "density" in label: + find_density = label.get("find_density", 1.0) + density_pred = model_pred["density"] + density_label = label["density"] + pref_d = pref_d * find_density + # mask out excluded grid points (the atomic model zeroes them + # and returns the mask) and reduce per frame, consistent with + # the other losses in the package, so that frames with larger + # grids do not dominate the batch + mask = model_pred.get("mask", None) + nframes = density_pred.shape[0] + if mask is None: + mask_f = torch.ones( + density_pred.shape[:2], + dtype=density_pred.dtype, + device=density_pred.device, + ) + else: + mask_f = mask.reshape(density_pred.shape[:2]).to(density_pred.dtype) + mask_sum = mask_f.sum(dim=-1).clamp(min=1.0) + residual = ( + density_label.reshape(nframes, -1) - density_pred.reshape(nframes, -1) + ) * mask_f + l2_density_loss = torch.square(residual).sum(dim=-1).div(mask_sum).mean() + rmse_d = l2_density_loss.sqrt() + more_loss["rmse_d"] = self.display_if_exist(rmse_d.detach(), find_density) + l1_density_loss = residual.abs().sum(dim=-1).div(mask_sum).mean() + mae_d = l1_density_loss + # minimise the squared error, consistent with every other loss + # in the package; the absolute error is only for display + loss += (pref_d * l2_density_loss).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_d"] = self.display_if_exist(mae_d.detach(), find_density) + elif not self.inference and "density" in model_pred: + # the density term is disabled (zero prefactors), but the graph + # must stay connected so that backward() does not fail + loss = loss + (model_pred["density"].sum() * 0.0).to( + GLOBAL_PT_FLOAT_PRECISION + ) + return model_pred, loss, more_loss + + def serialize(self) -> dict: + """Serialize the loss module. + + Returns + ------- + dict + The serialized loss module + """ + return { + "@class": "GridDensityLoss", + "@version": 1, + "starter_learning_rate": self.starter_learning_rate, + "start_pref_d": self.start_pref_d, + "limit_pref_d": self.limit_pref_d, + } + + @classmethod + def deserialize(cls, data: dict) -> "TaskLoss": + """Deserialize the loss module. + + Parameters + ---------- + data : dict + The serialized loss module + + Returns + ------- + TaskLoss + The deserialized loss module + """ + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) + + @property + def training_metric_names(self) -> tuple[str, ...]: + """Return density metrics displayed by averaged training output.""" + return ("rmse_d", "mae_d") if self.has_d else () + + @property + def label_requirement(self) -> list[DataRequirementItem]: + """Return data label requirements needed for this loss calculation. + + Only the density label is declared here; the grid is a model input + and is declared via ``get_additional_data_requirement``. + """ + label_requirement = [] + if self.has_d: + # the density label is the only supervision signal of this model: + # a missing file must abort training, not silently optimise nothing + label_requirement.append( + DataRequirementItem( + "density", + ndof=1, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + ) + return label_requirement diff --git a/deepmd/pt/model/atomic_model/__init__.py b/deepmd/pt/model/atomic_model/__init__.py index e343287cad..f1ec92eb84 100644 --- a/deepmd/pt/model/atomic_model/__init__.py +++ b/deepmd/pt/model/atomic_model/__init__.py @@ -17,6 +17,9 @@ from .base_atomic_model import ( BaseAtomicModel, ) +from .density_atomic_model import ( + DPDensityAtomicModel, +) from .dipole_atomic_model import ( DPDipoleAtomicModel, ) @@ -53,6 +56,7 @@ "BaseAtomicModel", "DPAtomicModel", "DPDOSAtomicModel", + "DPDensityAtomicModel", "DPDipoleAtomicModel", "DPEnergyAtomicModel", "DPPolarAtomicModel", diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index d18e780cd3..5497f33fe3 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -52,6 +52,7 @@ check_preset_out_bias, normalize_preset_out_bias, preset_out_bias_rows, + recover_nonfinite_preset_bias, remap_preset_out_bias, ) @@ -629,6 +630,7 @@ def change_out_bias( bias_adjust_mode: str = "change-by-statistic", model_forward: Callable[..., dict[str, torch.Tensor]] | None = None, observed_type: list[str] | None = None, + preset_out_bias: dict | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -654,6 +656,14 @@ def change_out_bias( ('set-by-statistic') or their stored bias ('change-by-statistic'); no statistics are computed, read or written for such an output, whose output std keeps its stored value. + preset_out_bias : dict, optional + An explicit downstream preset. Assigned outputs use this bias and + skip fitting; other outputs fit residuals against the stored bias. + The model preset is used when omitted, preserving training and + fine-tuning behavior. Pass an empty dict to fit from the data + without reapplying the model's stored preset. + Before residual fitting, non-finite stored cells covered by the + checkpoint preset are recovered; finite adjusted cells are retained. model_forward : Callable[..., dict[str, torch.Tensor]], optional Predictor of the complete atomic outputs used by 'change-by-statistic'. Defaults to the forward of this atomic model; a model that adds @@ -665,17 +675,34 @@ def change_out_bias( if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"): raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) change = bias_adjust_mode == "change-by-statistic" + preset = ( + normalize_preset_out_bias(preset_out_bias, self.type_map) + if preset_out_bias is not None + else self.preset_out_bias + ) distinguish_types = self.get_compute_stats_distinguish_types() - check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types) + check_preset_out_bias(preset, self.bias_keys, distinguish_types) sampled = sample_merged + # A per-call override must not propagate a corrupt checkpoint baseline. + if change and preset_out_bias is not None and self.preset_out_bias: + recovered = recover_nonfinite_preset_bias( + self.preset_out_bias, + self.type_map, + to_numpy_array(self.out_bias), + self.bias_keys, + [self.atomic_output_def()[kk].size for kk in self.bias_keys], + ) + self._store_out_stat( + {kk: to_torch_tensor(vv) for kk, vv in recovered.items()} + ) # === Step 1. Outputs fixed by the preset === fitted_keys = self.bias_keys - if self.preset_out_bias: + if preset: if observed_type is None: sampled = sample_merged() if callable(sample_merged) else sample_merged observed_type = collect_observed_types(sampled, self.type_map) rows = preset_out_bias_rows( - self.preset_out_bias, + preset, self.type_map, observed_type, to_numpy_array(self.out_bias), diff --git a/deepmd/pt/model/atomic_model/density_atomic_model.py b/deepmd/pt/model/atomic_model/density_atomic_model.py new file mode 100644 index 0000000000..3c79b714d3 --- /dev/null +++ b/deepmd/pt/model/atomic_model/density_atomic_model.py @@ -0,0 +1,687 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import numpy as np +import torch + +from deepmd.dpmodel.utils.env_mat_stat import ( + EnvMatStatSe, +) +from deepmd.pt.model.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.pt.model.descriptor.descriptor import ( + DescriptorBlock, +) +from deepmd.pt.model.task.density import ( + DensityFittingNet, +) +from deepmd.pt.utils.nlist import ( + build_directional_neighbor_list, + extend_input_and_build_neighbor_list, +) +from deepmd.pt.utils.region import ( + normalize_coord, +) +from deepmd.utils.model_preset import ( + PERIODIC_TABLE, +) +from deepmd.utils.path import ( + DPPath, +) + +from .dp_atomic_model import ( + DPAtomicModel, +) + +log = logging.getLogger(__name__) + + +class DPDensityAtomicModel(DPAtomicModel): + def __init__( + self, + descriptor: BaseDescriptor, + fitting: DensityFittingNet, + type_map: list[str], + **kwargs: Any, + ) -> None: + assert isinstance(fitting, DensityFittingNet) + super().__init__(descriptor, fitting, type_map, **kwargs) + self.rcut = self.descriptor.get_rcut() + self.rcut_smth = self.descriptor.get_rcut_smth() + if type_map and type_map[-1] in PERIODIC_TABLE: + raise ValueError( + "the last entry of type_map must be a reserved grid type " + "for a density model, but got the real element " + f"{type_map[-1]!r}; append a non-element placeholder " + '(e.g. "X") at the end of type_map' + ) + if self.descriptor.get_env_protection() == 0.0: + log.warning( + "The descriptor env_protection is 0.0; grid points coincident " + "with atoms would produce NaN densities. Setting it to 1e-6." + ) + self._set_descriptor_env_protection(1e-6) + self.sel = self.descriptor.get_sel() + self.nnei = self.descriptor.get_nsel() + + def forward_atomic( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + grid: torch.Tensor | None = None, + grid_type: torch.Tensor | None = None, + grid_nlist: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + return_atomic_feature: bool = False, + ) -> dict[str, torch.Tensor]: + """Return atomic prediction. + + Parameters + ---------- + extended_coord + coodinates in extended region + extended_atype + atomic type in extended region + nlist + neighbor list. nf x nloc x nsel + mapping + mapps the extended indices to local indices + fparam + frame parameter. nf x ndf + aparam + atomic parameter. nf x nloc x nda + comm_dict + Data needed for communication during parallel inference. + grid + Grid coordinates. nf x ngrid x 3. + grid_type + Reserved grid-point types. nf x ngrid. + grid_nlist + Directional grid-to-atom neighbor list. nf x ngrid x nsel. + charge_spin + Frame charge and spin conditions. nf x 2. + return_atomic_feature + Whether atomic features are requested. Density models do not + currently return them. + + Returns + ------- + result_dict + the result dict, defined by the `FittingOutputDef`. + + """ + del return_atomic_feature + nframes, _, _ = nlist.shape + if self.do_grad_r() or self.do_grad_c(): + extended_coord.requires_grad_(True) + assert mapping is not None + assert grid is not None + assert grid_type is not None + assert grid_nlist is not None + # Handle default chg_spin if the descriptor supports it, the same + # way DPAtomicModel.forward_atomic does + if self.add_chg_spin_ebd and charge_spin is None: + default_cs_tensor = self.descriptor.get_default_chg_spin() + if default_cs_tensor is not None: + default_cs_tensor = default_cs_tensor.to(device=extended_coord.device) + charge_spin = torch.tile(default_cs_tensor.unsqueeze(0), [nframes, 1]) + else: + raise ValueError( + "charge_spin is required when the descriptor enables " + "charge/spin embedding without default_chg_spin" + ) + _, ngrid, _ = grid_nlist.shape + # nb x (ngrid+nall) x 3 + merged_coord = torch.cat([grid, extended_coord], dim=1) + + grid_atype = torch.ones( + [nframes, ngrid], device=extended_atype.device, dtype=extended_atype.dtype + ) * (self.descriptor.get_ntypes() - 1) + # nb x (ngrid+nall) + merged_atype = torch.cat([grid_atype, extended_atype], dim=1) + + # nb x ngrid + grid_nlist_mask = grid_nlist >= 0 + shifted_grid_nlist = torch.where(grid_nlist_mask, grid_nlist + ngrid, -1) + # nb x all + nlist_mask = nlist >= 0 + shifted_nlist = torch.where(nlist_mask, nlist + ngrid, -1) + # nb x (ngrid+nall) + merged_nlist = torch.cat([shifted_grid_nlist, shifted_nlist], dim=1) + + # nb x ngrid, row i is the grid index i + grid_mapping = torch.arange( + ngrid, device=mapping.device, dtype=mapping.dtype + ).expand(nframes, -1) + # nb x (ngrid+nall) + merged_mapping = torch.cat([grid_mapping, mapping + ngrid], dim=1) + + # the descriptor evaluates the environment of every merged point; + # the atom rows are computed and discarded (only the grid rows feed + # the fitting net), which is inherent to evaluating the merged + # grid+atom system in one call + descriptor, rot_mat, g2, h2, _sw = self.descriptor( + merged_coord, + merged_atype, + merged_nlist, + mapping=merged_mapping, + comm_dict=comm_dict, + charge_spin=charge_spin if self.add_chg_spin_ebd else None, + ) + assert descriptor is not None + + # the fitting net must see the same grid type as the descriptor: + # the reserved last entry of the type map + grid_ftype = torch.ones( + [nframes, ngrid], device=grid_type.device, dtype=grid_type.dtype + ) * (self.descriptor.get_ntypes() - 1) + ret = self.fitting_net( + descriptor[:, :ngrid, :], + grid_ftype, + gr=rot_mat, + g2=g2, + h2=h2, + fparam=fparam, + ) + return ret + + def forward_common_atomic( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + grid: torch.Tensor | None = None, + grid_type: torch.Tensor | None = None, + grid_nlist: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Common interface for atomic inference. + + This method accept extended coordinates, extended atom typs, neighbor list, + and predict the atomic contribution of the fit property. + + Parameters + ---------- + extended_coord + extended coodinates, shape: nf x (nall x 3) + extended_atype + extended atom typs, shape: nf x nall + for a type < 0 indicating the atomic is virtual. + nlist + neighbor list, shape: nf x nloc x nsel + mapping + extended to local index mapping, shape: nf x nall + fparam + frame parameters, shape: nf x dim_fparam + aparam + atomic parameter, shape: nf x nloc x dim_aparam + comm_dict + The data needed for communication for parallel inference. + grid + Grid coordinates, shape: nf x ngrid x 3. + grid_type + Reserved grid-point types, shape: nf x ngrid. + grid_nlist + Directional grid-to-atom neighbor list, shape: nf x ngrid x nsel. + charge_spin + Frame charge and spin conditions, shape: nf x 2. + + Returns + ------- + ret_dict + dict of output atomic properties. + should implement the definition of `fitting_output_def`. + ret_dict["mask"] of shape nf x ngrid will be provided. + ret_dict["mask"][ff,ii] == 1 indicating the ii-th grid point of the ff-th frame is included. + ret_dict["mask"][ff,ii] == 0 indicating the ii-th grid point of the ff-th frame is excluded. + + """ + assert grid is not None + assert grid_type is not None + assert grid_nlist is not None + nframes, _, _ = nlist.shape + _, ngrid, _ = grid_nlist.shape + + if self.pair_excl is not None: + pair_mask = self.pair_excl(nlist, extended_atype) + # exclude neighbors in the nlist + nlist = torch.where(pair_mask == 1, nlist, -1) + # the directional grid-to-atom list: every center is the grid + # type, so exclusions involving the grid type drop the neighbor + reserved_type = self.descriptor.get_ntypes() - 1 + if any( + reserved_type in pair for pair in self.pair_excl.get_exclude_types() + ): + nall = extended_atype.shape[1] + nsel = grid_nlist.shape[-1] + virtual_type = self.pair_excl.ntypes * torch.ones( + [nframes, 1], + dtype=extended_atype.dtype, + device=extended_atype.device, + ) + ae = torch.cat([extended_atype, virtual_type], dim=-1) + index = torch.where(grid_nlist == -1, nall, grid_nlist).view( + nframes, ngrid * nsel + ) + type_j = torch.gather(ae, 1, index) + type_ij = reserved_type * (self.pair_excl.ntypes + 1) + type_j + grid_mask = self.pair_excl.type_mask[type_ij].view(nframes, ngrid, nsel) + grid_nlist = torch.where(grid_mask == 1, grid_nlist, -1) + + ext_atom_mask = self.make_atom_mask(extended_atype) + ret_dict = self.forward_atomic( + extended_coord, + torch.where(ext_atom_mask, extended_atype, 0), + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + comm_dict=comm_dict, + grid=grid, + grid_type=grid_type, + grid_nlist=grid_nlist, + charge_spin=charge_spin, + ) + ret_dict = self.apply_out_stat(ret_dict, grid_type) + + # nf x ngrid + grid_mask = torch.ones( + [nframes, ngrid], dtype=torch.int32, device=ext_atom_mask.device + ) + if self.atom_excl is not None: + grid_mask *= self.atom_excl(grid_type) + + for kk in ret_dict.keys(): + out_shape = ret_dict[kk].shape + out_shape2 = 1 + for ss in out_shape[2:]: + out_shape2 *= ss + ret_dict[kk] = ( + ret_dict[kk].reshape([out_shape[0], out_shape[1], out_shape2]) + * grid_mask[:, :, None] + ).view(out_shape) + ret_dict["mask"] = grid_mask + + return ret_dict + + def forward( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + raise NotImplementedError( + "DPDensityAtomicModel.forward requires grid, grid_type and grid_nlist; " + "use forward_common_atomic with explicit grid inputs instead." + ) + + def _get_forward_wrapper_func( + self, + ) -> Callable[..., dict[str, torch.Tensor]]: + """Get a forward wrapper of the atomic model for output bias calculation.""" + + def model_forward( + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + del spin + with torch.no_grad(): + ( + extended_coord, + extended_atype, + mapping, + nlist, + ) = extend_input_and_build_neighbor_list( + coord, + atype, + self.get_rcut(), + self.get_sel(), + mixed_types=self.mixed_types(), + box=box, + ) + assert grid is not None + if box is not None: + # same wrapping convention as forward_common: grid points + # outside the primary cell are periodically equivalent + grid = normalize_coord( + grid, box.to(grid.device).reshape(box.shape[0], 3, 3) + ) + grid_type = torch.full( + grid.shape[:-1], + self.descriptor.get_ntypes() - 1, + device=grid.device, + dtype=atype.dtype, + ) + grid_nlist = build_directional_neighbor_list( + grid, + grid_type, + extended_coord, + extended_atype, + self.get_rcut(), + self.get_sel(), + distinguish_types=(not self.mixed_types()), + ) + atomic_ret = self.forward_common_atomic( + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + grid=grid, + grid_type=grid_type, + grid_nlist=grid_nlist, + charge_spin=charge_spin, + ) + return {kk: vv.detach() for kk, vv in atomic_ret.items()} + + return model_forward + + def change_out_bias( + self, + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + bias_adjust_mode: str = "change-by-statistic", + model_forward: Callable[..., dict[str, torch.Tensor]] | None = None, + observed_type: list[str] | None = None, + ) -> None: + """Change the output bias according to the input data and the pretrained model. + + For density models, this operation is skipped because the output is + grid-based rather than atomic-based, and the standard atomic bias + adjustment (change-by-statistic / set-by-statistic) does not apply. + The fitting net will adapt to the target dataset through normal + gradient descent during training. + + This guard stays until output-stat support is designed for grid + models; tracked in https://github.com/deepmodeling/deepmd-kit/issues/6029. + """ + del ( + sample_merged, + stat_file_path, + bias_adjust_mode, + model_forward, + observed_type, + ) + log.warning("change_out_bias is not supported for density models; skipping.") + + def compute_or_load_stat( + self, + sampled_func: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, + preset_observed_type: list[str] | None = None, + ) -> None: + """Compute or load statistics, with real statistics for the grid type. + + The reserved grid type (``ntypes - 1``) has no real atoms in any + training system, so the standard pass yields zero samples for it + and falls back to placeholder statistics (mean 0, stddev 0.1), + which are then applied to exactly the rows that produce the + density output. When the sampled data provides grids, a second + pass is run with the grid points injected as pseudo-atoms of the + grid type, and only the grid-type row of the statistics is taken + from it; the real types keep their clean statistics from the + standard pass. + + The injected pass only runs when the grid-type statistics are + missing (no cache, or a cache whose grid-type row has zero + samples), so the stat-file fast path still applies; when a stat + file is used, the patched grid-type row is written back to disk so + that the file and the in-memory statistics agree. + """ + sampled = None + grid_stats = None + if self._grid_stat_missing(stat_file_path): + sampled = sampled_func() if callable(sampled_func) else sampled_func + if any("grid" in sample for sample in sampled): + try: + # run the wrapped sampler first so the injected pass sees + # the same pair_exclude_types as the standard pass + wrapped = self._make_wrapped_sampler(lambda: sampled) + self.descriptor.compute_input_stats( + self._inject_grid_samples(wrapped()) + ) + grid_stats = [ + ( + block["davg"][-1].detach().clone(), + block["dstd"][-1].detach().clone(), + ) + for block in self._descriptor_stat_blocks() + ] + except (TypeError, KeyError) as err: + log.warning( + "Cannot compute input statistics for the grid type (%s); " + "falling back to the descriptor defaults.", + err, + ) + else: + log.warning( + "No grid data in the sampled frames; the grid type gets the " + "descriptor's default input statistics." + ) + super().compute_or_load_stat( + (lambda: sampled) if sampled is not None else sampled_func, + stat_file_path, + compute_or_load_out_stat=compute_or_load_out_stat, + preset_observed_type=preset_observed_type, + ) + if grid_stats is not None: + for block, (grid_davg, grid_dstd) in zip( + self._descriptor_stat_blocks(), grid_stats, strict=True + ): + davg = block["davg"] + dstd = block["dstd"] + davg[-1] = grid_davg.to(device=davg.device, dtype=davg.dtype) + dstd[-1] = grid_dstd.to(device=dstd.device, dtype=dstd.dtype) + if ( + stat_file_path is not None + and getattr(stat_file_path, "mode", None) != "r" + ): + self._write_back_grid_stat(stat_file_path, grid_stats) + elif stat_file_path is not None: + log.warning( + "Grid-type statistics were patched in memory but were not " + "written back because the statistics cache is read-only." + ) + + def _descriptor_stat_blocks(self) -> list: + """Collect every descriptor block that carries davg/dstd statistics. + + Descriptors such as DPA-2 and hybrids carry several blocks, each with + its own statistics; the grid-type row must be patched in all of them. + The blocks are discovered by walking the module tree (via + ``named_modules``) and probing the statistics item protocol, so no + per-descriptor attribute list is needed. + """ + blocks = [] + for _name, module in self.descriptor.named_modules(): + # filter networks etc. also implement __getitem__ with unrelated + # semantics, so only descriptor-level modules are probed for the + # statistics item protocol (some descriptors, e.g. se_r, carry + # the statistics on the descriptor itself) + if not isinstance(module, (BaseDescriptor, DescriptorBlock)): + continue + try: + module["davg"] + module["dstd"] + except (TypeError, KeyError): + continue + blocks.append(module) + if blocks: + return blocks + try: + self.descriptor["davg"] + self.descriptor["dstd"] + return [self.descriptor] + except (TypeError, KeyError): + pass + raise KeyError("davg/dstd not accessible on this descriptor") + + def _set_descriptor_env_protection(self, value: float) -> None: + """Set env_protection on every descriptor block, hybrids included. + + Wrapper-level fields, if any, are derived from the blocks; the whole + module tree is walked so no per-descriptor attribute list is needed. + """ + for _name, module in self.descriptor.named_modules(): + if hasattr(module, "env_protection"): + module.env_protection = value + + def _stat_cache_root(self, stat_file_path: DPPath) -> DPPath: + """Apply the same type_map subdirectory as the parent's stat path.""" + if self.type_map is not None: + stat_file_path = stat_file_path / " ".join(self.type_map) + return stat_file_path + + def _grid_stat_missing(self, stat_file_path: DPPath | None) -> bool: + """Tell whether the grid-type input statistics still need computing. + + True when there is no complete stat cache, or when the cache is + complete but its grid-type row has zero samples (placeholder + statistics saved by an older run). + """ + if stat_file_path is None: + return True + try: + blocks = self._descriptor_stat_blocks() + except KeyError: + # no accessible statistics: let the injected pass soft-fail below + return True + stat_file_path = self._stat_cache_root(stat_file_path) + grid_type = self.descriptor.get_ntypes() - 1 + for block in blocks: + env_stat = EnvMatStatSe(block) + cache = stat_file_path / env_stat.get_hash() + keys = env_stat.get_stat_keys() + if not cache.is_dir() or any(not (cache / kk).is_file() for kk in keys): + return True + # Every component of the grid-type row must carry samples. In + # particular, a valid radial row must not hide a placeholder + # angular row left by an interrupted or legacy cache update. + suffix = f"_{grid_type}" + for key in keys: + if key.endswith(suffix) and float((cache / key).load_numpy()[0]) == 0.0: + return True + return False + + def _write_back_grid_stat( + self, + stat_file_path: DPPath, + grid_stats: list[tuple[torch.Tensor, torch.Tensor]], + ) -> None: + """Write the patched grid-type row back to the stat cache files. + + The cache stores raw sums (number, sum, squared_sum) per stat key; + a single sample reproduces the patched mean and stddev exactly. + """ + grid_type = self.descriptor.get_ntypes() - 1 + stat_file_path = self._stat_cache_root(stat_file_path) + for block, (grid_davg, grid_dstd) in zip( + self._descriptor_stat_blocks(), grid_stats, strict=True + ): + env_stat = EnvMatStatSe(block) + cache = stat_file_path / env_stat.get_hash() + if not cache.is_dir(): + continue + davg_np = grid_davg.detach().cpu().numpy() + dstd_np = grid_dstd.detach().cpu().numpy() + mean_r = float(davg_np[0, 0]) + std_r = float(dstd_np[0, 0]) + (cache / f"r_{grid_type}").save_numpy( + np.array([1.0, mean_r, std_r**2 + mean_r**2]) + ) + if dstd_np.shape[-1] == 4: + std_a = float(dstd_np[0, 1]) + (cache / f"a_{grid_type}").save_numpy(np.array([1.0, 0.0, std_a**2])) + + def _inject_grid_samples(self, sampled: list[dict]) -> list[dict]: + """Append the grid points of each sample as pseudo-atoms of the + reserved grid type, so the descriptor input statistics get real + samples for it. Neighbor lists are rebuilt by the stat machinery, + so only ``coord`` and ``atype`` need to be extended. + + Grid-grid pairs are excluded from the injected neighbor lists so + that the grid-type statistics are computed over grid-to-atom + pairs, matching the directional neighbor list of the forward pass + (a no-op for per-type-sel descriptors with ``sel[-1] == 0``, but + required for mixed-type descriptors with a scalar sel). + """ + ntypes = self.descriptor.get_ntypes() + grid_type = ntypes - 1 + injected = [] + for sample in sampled: + grid = sample.get("grid") + if grid is None: + injected.append(sample) + continue + sample = dict(sample) + coord = sample["coord"] + atype = sample["atype"] + nframes = atype.shape[0] + gg = grid.reshape(nframes, -1, 3).to(coord.dtype) + ngrid = gg.shape[1] + sample["coord"] = torch.cat( + [coord.reshape(nframes, -1, 3), gg], dim=1 + ).reshape(nframes, -1) + sample["atype"] = torch.cat( + [ + atype, + torch.full( + (nframes, ngrid), + grid_type, + dtype=atype.dtype, + device=atype.device, + ), + ], + dim=1, + ) + pair_exclude_types = [ + tuple(pair) for pair in sample.get("pair_exclude_types", []) + ] + if (grid_type, grid_type) not in pair_exclude_types: + pair_exclude_types.append((grid_type, grid_type)) + sample["pair_exclude_types"] = pair_exclude_types + injected.append(sample) + return injected + + def compute_or_load_out_stat( + self, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the output statistics for the fitting net. + + Not implemented for density models: the output lives on grid points, + so per-type atomic output statistics do not apply. This guard stays + until output-stat support is designed for grid models; tracked in + https://github.com/deepmodeling/deepmd-kit/issues/6029. + """ + log.warning("Not implemented yet for density out stat!") diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index d2e0a09e96..e3ea613b44 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -49,6 +49,9 @@ env_protection_descriptors, ) +from .density_model import ( + GridDensityModel, +) from .dipole_model import ( DipoleModel, ) @@ -99,15 +102,21 @@ from .sezm_property_model import ( SeZMPropertyModel, ) -from .sezm_spin_model import ( - SeZMSpinModel, -) from .spin_model import ( SpinEnergyModel, SpinModel, ) +def _reject_long_range(model_params: dict) -> None: + """Reject a long-range module, which only the pt_expt backend serves.""" + if model_params.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) + + def _get_standard_model_components(model_params: dict, ntypes: int) -> tuple: if "type_embedding" in model_params: raise ValueError( @@ -208,6 +217,7 @@ def get_linear_model(model_params: dict) -> BaseModel: "sub-model: add an `inner_potential` sub-model to the " "composition instead." ) + _reject_long_range(sub) if any( sub.get("type") == "inner_potential" for sub in model_params.get("models", []) ): @@ -373,6 +383,7 @@ def _get_bridged_linear_model(model_params: dict) -> BaseModel: def get_zbl_model(model_params: dict) -> DPZBLModel: + _reject_long_range(model_params) model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) descriptor, fitting, _ = _get_standard_model_components(model_params, ntypes) @@ -414,6 +425,7 @@ def get_standard_model(model_params: dict) -> BaseModel: "expands the flag), or spell the composition explicitly with " '`type: "linear_ener"` and an `inner_potential` sub-model.' ) + _reject_long_range(model_params) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) @@ -433,6 +445,8 @@ def get_standard_model(model_params: dict) -> BaseModel: modelcls = DOSModel elif fitting_net_type in ["ener", "direct_force_ener"]: modelcls = EnergyModel + elif fitting_net_type == "density": + modelcls = GridDensityModel elif fitting_net_type == "property": modelcls = PropertyModel elif fitting_net_type == "population": @@ -593,29 +607,14 @@ def get_sezm_model(model_params: dict) -> BaseModel: def get_sezm_spin_model(model_params: dict) -> BaseModel: - """Dispatch a DPA4/SeZM spin model by ``model.spin.scheme``. - - ``deepspin`` selects the classical virtual-atom representation; ``native`` - injects the spin into the descriptor as an equivariant feature. A missing - ``scheme`` falls back to ``deepspin``. - """ - scheme = str(model_params.get("spin", {}).get("scheme", "deepspin")).lower() - if scheme == "native": - return _get_sezm_native_spin_model(model_params) - if scheme == "deepspin": - return _get_sezm_virtual_spin_model(model_params) - raise ValueError( - f"Unknown spin scheme '{scheme}' for DPA4/SeZM; use 'native' or 'deepspin'." - ) - - -def _get_sezm_native_spin_model(model_params: dict) -> BaseModel: - """Build the native (virtual-atom-free) spin DPA4/SeZM model. + """Build the native spin DPA4/SeZM model. The spin vector enters the descriptor as an equivariant feature, so the type map, neighbor selection and type count stay at the real-system sizes; only the descriptor gains the spin-embedding branch. """ + if str(model_params.get("spin", {}).get("scheme", "")).lower() != "native": + raise ValueError("DPA4/SeZM spin models require model.spin.scheme='native'.") model_params_old = model_params model_params = copy.deepcopy(model_params) model_params.setdefault("descriptor", {}) @@ -624,8 +623,6 @@ def _get_sezm_native_spin_model(model_params: dict) -> BaseModel: _normalize_spin_use_spin(model_params) use_spin = [bool(flag) for flag in model_params["spin"]["use_spin"]] - # ``virtual_scale`` is a virtual-atom geometric device; the native scheme - # only needs ``use_spin`` for masking, so it stays unset here. spin = Spin( use_spin=use_spin, allow_missing_label=model_params["spin"].get("allow_missing_label", False), @@ -641,7 +638,7 @@ def _get_sezm_native_spin_model(model_params: dict) -> BaseModel: if pair_exclude_types: model_params["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types) - # === Bridging parameters (no virtual atoms, so ZBL needs no masking) === + # === Bridging parameters === bridging_method = str(model_params.get("bridging_method", "none")).upper() bridging_window = _bridging_window_options(model_params) if bridging_method != "NONE": @@ -685,77 +682,6 @@ def _get_sezm_native_spin_model(model_params: dict) -> BaseModel: return model -def _get_sezm_virtual_spin_model(model_params: dict) -> BaseModel: - model_params_old = model_params - model_params = copy.deepcopy(model_params) - model_params.setdefault("descriptor", {}) - model_params.setdefault("fitting_net", {}) - model_params["descriptor"].setdefault("type", "dpa4") - _normalize_spin_use_spin(model_params) - real_sel = model_params["descriptor"].get("sel", 120) - real_sel_list = [int(real_sel)] if isinstance(real_sel, int) else list(real_sel) - real_nsel = int(sum(real_sel_list)) - model_params["descriptor"]["sel"] = [2 * real_nsel + 1] - - spin = Spin( - use_spin=model_params["spin"]["use_spin"], - virtual_scale=model_params["spin"]["virtual_scale"], - allow_missing_label=model_params["spin"].get("allow_missing_label", False), - ) - real_pair_exclude_types = _reconcile_sezm_pair_exclude_types(model_params) - model_params["type_map"] += [item + "_spin" for item in model_params["type_map"]] - pair_exclude_types = spin.get_pair_exclude_types( - exclude_types=real_pair_exclude_types or None - ) - model_params["pair_exclude_types"] = pair_exclude_types - model_params["descriptor"]["exclude_types"] = pair_exclude_types - atom_exclude_types = spin.get_atom_exclude_types( - exclude_types=model_params.get("atom_exclude_types", None) - ) - model_params["atom_exclude_types"] = atom_exclude_types - - ntypes = len(model_params["type_map"]) - model_params["descriptor"]["ntypes"] = ntypes - model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"]) - - descriptor = BaseDescriptor(**model_params["descriptor"]) - - fitting_net = copy.deepcopy(model_params["fitting_net"]) - fitting_net_type = fitting_net.get("type", "dpa4_ener") - if fitting_net_type not in ("dpa4_ener", "sezm_ener"): - raise ValueError( - "Spin DPA4/SeZM currently supports only `dpa4_ener` or `sezm_ener` " - f"fitting, but got `{fitting_net_type}`." - ) - fitting_net.pop("type", None) - fitting_net["ntypes"] = descriptor.get_ntypes() - fitting_net["type_map"] = copy.deepcopy(model_params["type_map"]) - fitting_net["mixed_types"] = descriptor.mixed_types() - fitting_net["dim_descrpt"] = descriptor.get_dim_out() - fitting = SeZMEnergyFittingNet(**fitting_net) - preset_out_bias = model_params.get("preset_out_bias") - data_stat_protect = model_params.get("data_stat_protect", 1e-2) - use_compile = bool(model_params.get("use_compile", False)) - enable_tf32 = bool(model_params.get("enable_tf32", True)) - - model = SeZMSpinModel( - descriptor=descriptor, - fitting=fitting, - type_map=model_params["type_map"], - atom_exclude_types=atom_exclude_types, - pair_exclude_types=pair_exclude_types, - preset_out_bias=preset_out_bias, - data_stat_protect=data_stat_protect, - use_compile=use_compile, - enable_tf32=enable_tf32, - bridging_method=str(model_params.get("bridging_method", "none")), - real_sel=real_sel_list, - spin=spin, - ) - model.model_def_script = json.dumps(model_params_old) - return model - - def get_model(model_params: dict) -> Any: model_params = expand_bridging_method(model_params) model_type = model_params.get("type", "standard") @@ -790,7 +716,6 @@ def get_model(model_params: dict) -> Any: "SeZMModel", "SeZMNativeSpinModel", "SeZMPropertyModel", - "SeZMSpinModel", "SpinEnergyModel", "SpinModel", "get_model", diff --git a/deepmd/pt/model/model/density_model.py b/deepmd/pt/model/model/density_model.py new file mode 100644 index 0000000000..a319bf6cf5 --- /dev/null +++ b/deepmd/pt/model/model/density_model.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from copy import ( + deepcopy, +) +from typing import ( + Any, +) + +import torch + +from deepmd.pt.model.atomic_model import ( + DPDensityAtomicModel, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) + +from .dp_model import ( + DPModelCommon, +) +from .make_density_model import ( + make_density_model, +) + +DPDensityModel_ = make_density_model(DPDensityAtomicModel) + + +@BaseModel.register("density") +@BaseModel.register("grid_density") +class GridDensityModel(DPModelCommon, DPDensityModel_): + model_type = "grid_density" + + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPDensityModel_.__init__(self, *args, **kwargs) + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "density": deepcopy(out_def_data["density"]), + } + if "mask" in out_def_data: + output_def["mask"] = deepcopy(out_def_data["mask"]) + return output_def + + @torch.jit.export + def has_grid(self) -> bool: + """Returns whether it has grid input and output.""" + return True + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + assert grid is not None + model_ret = self.forward_common( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + grid=grid, + ) + model_predict = {} + model_predict["density"] = model_ret["density"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + @torch.jit.export + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, + ) -> None: + raise NotImplementedError("forward_lower is not supported for density models.") diff --git a/deepmd/pt/model/model/make_density_model.py b/deepmd/pt/model/model/make_density_model.py new file mode 100644 index 0000000000..b41daf0180 --- /dev/null +++ b/deepmd/pt/model/model/make_density_model.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + Any, +) + +import torch + +from deepmd.pt.model.atomic_model.base_atomic_model import ( + BaseAtomicModel, +) +from deepmd.pt.model.model.make_model import ( + make_model, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) +from deepmd.pt.model.model.transform_output import ( + communicate_extended_output, + fit_output_to_model_output, +) +from deepmd.pt.utils.nlist import ( + build_directional_neighbor_list, + extend_input_and_build_neighbor_list, +) +from deepmd.pt.utils.region import ( + normalize_coord, +) + +log = logging.getLogger(__name__) + + +def make_density_model(T_AtomicModel: type[BaseAtomicModel]) -> type[BaseModel]: + """Make a density model as a derived class of an atomic model. + + The model predicts a scalar density on grid points instead of atomic + energies. It is built by subclassing the standard model made by + `make_model` and overriding only the interfaces that differ: + + 1. the `forward_common_lower`, that takes extended coordinates, atypes, + neighbor list and the grid (with its type and directional neighbor list), + and outputs the density on the grid points; + + 2. the `forward_common`, that takes coordinates, atypes, cell and grid, + and predicts the density on the grid points. + + Parameters + ---------- + T_AtomicModel + The atomic model. + + Returns + ------- + CM + The model. + + """ + + class CM(make_model(T_AtomicModel)): + def forward_common( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + coord_corr_for_virial: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Return model prediction. + + The base parameters keep the base positional order (including + ``coord_corr_for_virial``, which is accepted and ignored since + density has no virial correction); the density-specific + ``grid`` is appended at the end and is required for this model. + + Parameters + ---------- + coord + The coordinates of the atoms. + shape: nf x (nloc x 3) + atype + The type of atoms. shape: nf x nloc + box + The simulation box. shape: nf x 9 + fparam + frame parameter. nf x ndf + aparam + atomic parameter. nf x nloc x nda + do_atomic_virial + If calculate the atomic virial. + grid + The coordinates of the grids. + shape: nf x (ngrid x 3) + + Returns + ------- + ret_dict + The result dict of type dict[str,torch.Tensor]. + The keys are defined by the `ModelOutputDef`. + + """ + del coord_corr_for_virial + assert grid is not None + cc, gg, bb, fp, ap, input_prec = self._input_type_cast( + coord, box=box, fparam=fparam, aparam=aparam, grid=grid + ) + del coord, grid, box, fparam, aparam + gg = gg.view(gg.shape[0], -1, 3) + if gg.shape[1] == 0: + raise ValueError("grid must contain at least one point per frame") + if bb is not None: + # wrap grid points into the primary cell: the atomic + # coordinates are normalized and extended with ghosts, so + # periodically equivalent grids outside the cell would + # otherwise lose their neighbors + gg = normalize_coord(gg, bb.to(gg.device).view(bb.shape[0], 3, 3)) + ( + extended_coord, + extended_atype, + mapping, + nlist, + ) = extend_input_and_build_neighbor_list( + cc, + atype, + self.get_rcut(), + self.get_sel(), + mixed_types=self.mixed_types(), + box=bb, + ) + grid_type = torch.full( + (gg.shape[0], gg.shape[1]), + self.atomic_model.descriptor.get_ntypes() - 1, + device=gg.device, + dtype=atype.dtype, + ) + grid_nlist = build_directional_neighbor_list( + gg, + grid_type, + extended_coord, + extended_atype, + self.get_rcut(), + self.get_sel(), + distinguish_types=(not self.mixed_types()), + ) + model_predict_lower = self.forward_common_lower( + extended_coord, + extended_atype, + nlist, + grid=gg, + grid_type=grid_type, + grid_nlist=grid_nlist, + mapping=mapping, + do_atomic_virial=do_atomic_virial, + fparam=fp, + aparam=ap, + charge_spin=charge_spin, + ) + model_predict = communicate_extended_output( + model_predict_lower, + self.model_output_def(), + mapping, + do_atomic_virial=do_atomic_virial, + ) + model_predict = self._output_type_cast(model_predict, input_prec) + return model_predict + + def forward_common_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + comm_dict: dict[str, torch.Tensor] | None = None, + extra_nlist_sort: bool = False, + extended_coord_corr: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + grid_type: torch.Tensor | None = None, + grid_nlist: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Return model prediction. Lower interface that takes + extended atomic coordinates and types, nlist, mapping, and the + grid inputs as input, and returns the predictions on the grid + points. The predictions are not reduced. + + The base parameters keep the base positional order (including + ``extended_coord_corr``, which is accepted and ignored); the + density-specific grid parameters are appended at the end and + are required for this model. + + Parameters + ---------- + extended_coord + coodinates in extended region. nf x (nall x 3) + extended_atype + atomic type in extended region. nf x nall + nlist + neighbor list. nf x nloc x nsel. + grid + grid coordinates. nf x ngrid x 3 + grid_type + type of the grid points. nf x ngrid + grid_nlist + directional neighbor list from grid points to atoms. + nf x ngrid x nsel + mapping + mapps the extended indices to local indices. nf x nall. + fparam + frame parameter. nf x ndf + aparam + atomic parameter. nf x nloc x nda + do_atomic_virial + whether calculate atomic virial. + comm_dict + The data needed for communication for parallel inference. + extra_nlist_sort + whether to forcibly sort the nlist. + + Returns + ------- + result_dict + the result dict, defined by the `FittingOutputDef`. + + """ + del extended_coord_corr + nframes, nall = extended_atype.shape[:2] + extended_coord = extended_coord.view(nframes, -1, 3) + nlist = self.format_nlist( + extended_coord, extended_atype, nlist, extra_nlist_sort=extra_nlist_sort + ) + assert grid is not None + assert grid_type is not None + assert grid_nlist is not None + if grid.shape[1] == 0: + raise ValueError("grid must contain at least one point per frame") + cc_ext, gg, _, fp, ap, input_prec = self._input_type_cast( + extended_coord, grid=grid, fparam=fparam, aparam=aparam + ) + del extended_coord, grid, fparam, aparam + atomic_ret = self.atomic_model.forward_common_atomic( + cc_ext, + extended_atype, + nlist, + mapping=mapping, + fparam=fp, + aparam=ap, + comm_dict=comm_dict, + grid=gg, + grid_type=grid_type, + grid_nlist=grid_nlist, + charge_spin=charge_spin, + ) + model_predict = fit_output_to_model_output( + atomic_ret, + self.atomic_output_def(), + cc_ext, + do_atomic_virial=do_atomic_virial, + create_graph=self.training, + ) + model_predict = self._output_type_cast(model_predict, input_prec) + return model_predict + + def _input_type_cast( + self, + coord: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + str, + ]: + """Cast the input data to global float type.""" + input_prec = self.reverse_precision_dict[coord.dtype] + # dtype mismatch warnings are not emitted here: type checking + # would not pass jit, so inputs are converted to coord prec anyway. + _lst: list[torch.Tensor | None] = [ + vv.to(coord.dtype) if vv is not None else None + for vv in [grid, box, fparam, aparam] + ] + grid, box, fparam, aparam = _lst + assert grid is not None + if ( + input_prec + == self.reverse_precision_dict[self.global_pt_float_precision] + ): + return coord, grid, box, fparam, aparam, input_prec + else: + pp = self.global_pt_float_precision + return ( + coord.to(pp), + grid.to(pp), + box.to(pp) if box is not None else None, + fparam.to(pp) if fparam is not None else None, + aparam.to(pp) if aparam is not None else None, + input_prec, + ) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + # directly call the forward_common method when no specific transform rule + assert grid is not None + return self.forward_common( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + grid=grid, + ) + + @torch.jit.export + def forward_embedding( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + # the base implementation would call the grid-aware caster without + # a grid; density models do not support embedding extraction + raise NotImplementedError( + "forward_embedding is not supported for density models." + ) + + def change_out_bias( + self, + merged: Any, + bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, + ) -> None: + """Change the output bias according to the input data. + + Not supported for density models: the output is defined on grid + points rather than atoms, so the standard bias adjustment (which + would run a grid-less forward through the stat wrapper) does not + apply. Overridden at the model level because the default + ``change-by-statistic`` mode never reaches the atomic-level + no-op. + """ + del merged, bias_adjust_mode, preset_out_bias + log.warning( + "change_out_bias is not supported for density models; skipping." + ) + + return CM diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 185122349d..9b83dc8476 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -335,6 +335,7 @@ def change_out_bias( self, merged: Any, bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. @@ -355,10 +356,16 @@ def change_out_bias( 'change-by-statistic' : perform predictions on labels of target dataset, and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. + preset_out_bias : dict, optional + Per-call preset override. Omit to use the configured preset; + pass an empty dict to fit without fixing preset outputs. + Non-finite stored cells covered by the checkpoint preset are + recovered before residual fitting; finite values are retained. """ self.atomic_model.change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, model_forward=self.predict_atomic_outputs_for_stat, ) diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index b1fdf148db..75b2d14a4c 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -89,7 +89,7 @@ '-- communicate_extended_output Subsequent batches look up the cached callable at the same -``(training, has_coord_corr)`` slot of ``compiled_core_compute_cache``. +``training`` slot of ``compiled_core_compute_cache``. Each slot is retained independently, so train <-> eval toggles around every ``disp_freq`` / full-validation checkpoint reuse the other slot's compile product instead of evicting it (NOTE 7). @@ -281,19 +281,22 @@ Workaround for PyTorch <=2.11 bugs pytorch/pytorch#174379, #178080, #179494. All three manifest only under data-dependent symbolic shapes -- exactly our edge count. +* ``inplace_buffers=False`` + Triton 3.7.1 can compute a reduction wrongly when the kernel stores + its result over an input that it loads in the same loop and the + tuned configuration replicates the block across warps; separate + output buffers leave no kernel that stores over its own inputs. NOTE 7 -- Multi-slot compile cache key -------------------------------------- -The key is ``(training, has_coord_corr)``: +The key is ``training``: * ``self.training`` switches ``create_graph`` in ``edge_energy_deriv`` -- it toggles the entire second-derivative branch on or off. -* ``has_coord_corr`` selects the spin-virial correction branch, changing the - compiled callable arity from six tensor inputs to seven. No single compiled graph can serve both variants, so the cache is a -``dict[tuple[bool, bool], Callable]`` named +``dict[bool, Callable]`` named ``compiled_core_compute_cache``. A single-slot cache would have to evict on every flip, which turns the normal training-loop pattern -- ``train -> eval at every disp_freq -> train`` @@ -511,17 +514,12 @@ strip_saved_tensor_detach, trace_pad_dim, ) -from deepmd.pt.utils.nv_nlist import ( - NvNeighborList, - is_nv_available, +from deepmd.pt_expt.utils.auto_neighbor_list import ( + AdaptiveNeighborList, ) from deepmd.pt_expt.utils.edge_schema import ( edge_schema_from_extended, ) -from deepmd.pt_expt.utils.vesin_neighbor_list import ( - VesinNeighborList, - is_vesin_torch_available, -) from deepmd.utils.version import ( check_version_compatibility, ) @@ -531,49 +529,6 @@ SeZMModel_ = make_model(SeZMAtomicModel) -def _select_neighbor_builder(nf: int, device: torch.device) -> NeighborList: - """Select the O(N) neighbor builder for the given batch shape and device. - - A single CPU frame uses ``vesin``: its lightweight cell list has no - per-frame batching to amortize and avoids the ``nvalchemiops`` (Warp) - launch overhead. Every other case -- any CUDA input or any multi-frame - batch -- uses ``nvalchemiops``, whose batched kernel amortizes the launch - cost across frames. Both builders keep every neighbor within ``rcut``, so - the resulting edge set never depends on ``sel``. ``vesin`` also serves as - the fallback when ``nvalchemiops`` is absent. - - Parameters - ---------- - nf : int - Number of frames in the batch. - device : torch.device - Device of the input coordinates. - - Returns - ------- - NeighborList - A builder instance whose ``build`` method produces the neighbor list. - - Raises - ------ - RuntimeError - If neither ``nvalchemiops`` nor ``vesin`` is importable. - """ - if device.type == "cpu" and nf == 1 and is_vesin_torch_available(): - return VesinNeighborList() - if is_nv_available(): - return NvNeighborList() - if is_vesin_torch_available(): - return VesinNeighborList() - raise RuntimeError( - "SeZM neighbor-list construction requires a neighbor-list backend, but " - "neither 'vesin' nor 'nvalchemiops' is importable. Install one of:\n" - " pip install 'vesin[torch]' # portable CPU/CUDA cell list\n" - " pip install nvalchemi-toolkit-ops # batched CUDA kernels\n" - "'vesin' also ships with the PyTorch extra: pip install 'deepmd-kit[torch]'." - ) - - # Apply the process-global PyTorch workarounds the compile pipeline relies on # (autotune log suppression, DDP optimiser, and the 2.12 divisibility repair) # once, before the first compilation in this run. @@ -582,7 +537,7 @@ def _select_neighbor_builder(nf: int, device: torch.device) -> NeighborList: # --------------------------------------------------------------------------- # Multi-task compile sharing # --------------------------------------------------------------------------- -# Maps (structure_key..., training, has_coord_corr) to the compiled callable. +# Maps (structure_key..., training) to the compiled callable. # Tasks whose descriptor and fitting parameters share the same Python-object # identity after ``share_params(level=0)`` reuse one compiled graph, avoiding # N x compile-cache growth and duplicated DDP graph boundaries. @@ -719,16 +674,17 @@ def __init__( self.redu_prec = env.GLOBAL_PT_ENER_FLOAT_PRECISION self.use_compile = bool(use_compile) self.enable_tf32 = bool(enable_tf32) + self._inference_neighbor_list: NeighborList | None = None # LoRA injection happens in Trainer.__init__ after pre-trained state is loaded. self.lora_config: dict[str, Any] | None = None if lora is None else dict(lora) self._dens_compiled = False self._core_compute_pending_compile_t0: float | None = None - self._core_compute_pending_compile_key: tuple[bool, bool] | None = None + self._core_compute_pending_compile_key: bool | None = None self._dens_pending_compile_t0: float | None = None # Store compiled callables outside the nn.Module tree so that # FSDP2 / DDP do not shard or sync its duplicated parameters. # ``compiled_core_compute_cache`` is keyed on - # ``(training, has_coord_corr)`` so every graph topology has its own + # ``training`` so every graph topology has its own # slot; flipping between train and eval for validation -- regular, # full, or EMA full -- therefore reuses cached compile products # instead of evicting the other mode. @@ -1078,7 +1034,6 @@ def forward_common_lower( aparam: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, extended_atype: torch.Tensor | None = None, - extended_coord_corr: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, input_prec: torch.dtype | None = None, @@ -1120,12 +1075,6 @@ def forward_common_lower( device=coord.device, dtype=torch.long ) edge_mask = edge_mask.to(device=coord.device, dtype=torch.bool) - if extended_coord_corr is not None: - extended_coord_corr = extended_coord_corr.to( - device=coord.device, dtype=coord.dtype - ) - if extended_coord_corr.ndim == 2: - extended_coord_corr = extended_coord_corr.reshape(atype.shape[0], -1, 3) if spin is not None: spin = spin.to(device=coord.device, dtype=coord.dtype) nf = atype.shape[0] @@ -1157,10 +1106,9 @@ def forward_common_lower( dtype=coord.dtype, device=coord.device, ) - has_coord_corr = extended_coord_corr is not None if embedding_only: # Eval-only graph: a single slot, with no - # (training, coord_corr) key and no cross-task sharing. + # training key and no cross-task sharing. if self.compiled_embedding is None: self.trace_and_compile( coord, @@ -1178,7 +1126,7 @@ def forward_common_lower( compiled_core_compute = self.compiled_embedding task_buf_order = self._embedding_task_buf_order else: - cache_key = (bool(self.training), has_coord_corr) + cache_key = bool(self.training) if cache_key not in self.compiled_core_compute_cache: self.trace_and_compile( coord, @@ -1190,7 +1138,6 @@ def forward_common_lower( fp, ap, charge_spin, - extended_coord_corr=extended_coord_corr, spin=spin, ) compiled_core_compute = self.compiled_core_compute_cache[cache_key] @@ -1213,19 +1160,6 @@ def forward_common_lower( spin, *task_buf_vals, ) - elif extended_coord_corr is None: - model_predict = compiled_core_compute( - coord, - atype, - edge_index, - edge_vec, - edge_scatter_index, - edge_mask, - fp, - ap, - charge_spin, - *task_buf_vals, - ) else: model_predict = compiled_core_compute( coord, @@ -1237,7 +1171,6 @@ def forward_common_lower( fp, ap, charge_spin, - extended_coord_corr, *task_buf_vals, ) if ( @@ -1247,9 +1180,8 @@ def forward_common_lower( if torch.cuda.is_available(): torch.cuda.synchronize() log.info( - "SeZM: finished compiling (mode=%s, coord_corr=%s) in %.2fs", + "SeZM: finished compiling (mode=%s) in %.2fs", "train" if self.training else "eval", - has_coord_corr, time.perf_counter() - self._core_compute_pending_compile_t0, extra={"rank_scope": "all"}, ) @@ -1269,7 +1201,6 @@ def forward_common_lower( charge_spin=charge_spin, comm_dict=comm_dict, extended_atype=extended_atype, - extended_coord_corr=extended_coord_corr, spin=spin, embedding_only=embedding_only, atomic_output_only=atomic_output_only, @@ -1415,7 +1346,6 @@ def core_compute( charge_spin: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, extended_atype: torch.Tensor | None = None, - extended_coord_corr: torch.Tensor | None = None, spin: torch.Tensor | None = None, embedding_only: bool = False, conservative: bool = True, @@ -1457,9 +1387,6 @@ def core_compute( Extended atom types with shape ``(nf, nall)``. Required when ``comm_dict`` is provided so the descriptor can embed ghost types; unused otherwise. - extended_coord_corr - Coordinates correction for virial with shape ``(nf, nscatter, 3)`` or - ``None``. spin Optional per-atom spin vectors with shape ``(nf, nloc, 3)`` for the native spin scheme. When provided on the conservative path, the @@ -1648,7 +1575,6 @@ def core_compute( coord, create_graph=False, mask=fit_ret["mask"], - extended_coord_corr=extended_coord_corr, ) # === Step 6. Force / virial via edge-force scatter === @@ -1672,7 +1598,6 @@ def core_compute( nf, nscatter, create_graph=self.training, - extended_coord_corr=extended_coord_corr, spin_leaf=spin, ) ) @@ -1921,7 +1846,6 @@ def trace_and_compile( fp: torch.Tensor, ap: torch.Tensor, charge_spin: torch.Tensor, - extended_coord_corr: torch.Tensor | None = None, spin: torch.Tensor | None = None, embedding_only: bool = False, ) -> None: @@ -1943,7 +1867,6 @@ def trace_and_compile( ) mode = "train" if self.training else "eval" - has_coord_corr = extended_coord_corr is not None _compile_t0 = time.perf_counter() # --- Check module-level shared cache first --- @@ -1951,24 +1874,22 @@ def trace_and_compile( # should share one compiled graph. If a sibling task already compiled, # populate this instance's per-instance caches and return immediately. structure_key = _sezm_structure_key(self) - cache_key = (bool(self.training), has_coord_corr) - full_cache_key = structure_key + cache_key + cache_key = bool(self.training) + full_cache_key = (*structure_key, cache_key) if not embedding_only and full_cache_key in _SEZM_COMPILE_CACHE: self.compiled_core_compute_cache[cache_key] = _SEZM_COMPILE_CACHE[ full_cache_key ] self._task_buf_order_cache[cache_key] = _SEZM_TASK_BUF_ORDER[structure_key] log.debug( - "SeZM: reusing shared compiled graph (mode=%s, coord_corr=%s)", + "SeZM: reusing shared compiled graph (mode=%s)", mode, - has_coord_corr, ) return log.info( - "SeZM: start tracing and compiling (mode=%s, coord_corr=%s)", + "SeZM: start tracing and compiling (mode=%s)", mode, - has_coord_corr, extra={"rank_scope": "all"}, ) @@ -2085,40 +2006,9 @@ def compute_fn( # type: ignore[misc] finally: _restore_task_bufs(_saved) - elif extended_coord_corr is None: - - def compute_fn( - coord: torch.Tensor, - atype: torch.Tensor, - edge_index: torch.Tensor, - edge_vec: torch.Tensor, - edge_scatter_index: torch.Tensor, - edge_mask: torch.Tensor, - fp: torch.Tensor, - ap: torch.Tensor, - charge_spin: torch.Tensor, - *task_buf_vals: torch.Tensor, - ) -> dict[str, torch.Tensor]: - _saved = _patch_task_bufs(task_buf_vals) - try: - return self.core_compute( - _prepare_coord_for_trace(coord), - atype, - edge_index, - edge_vec, - edge_scatter_index, - edge_mask, - fparam=fp, - aparam=ap, - charge_spin=charge_spin, - embedding_only=embedding_only, - ) - finally: - _restore_task_bufs(_saved) - else: - def compute_fn( # type: ignore[misc] + def compute_fn( coord: torch.Tensor, atype: torch.Tensor, edge_index: torch.Tensor, @@ -2128,12 +2018,8 @@ def compute_fn( # type: ignore[misc] fp: torch.Tensor, ap: torch.Tensor, charge_spin: torch.Tensor, - extended_coord_corr: torch.Tensor, *task_buf_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: - # NOTE: Spin virial uses a coordinate correction derived from the - # virtual-atom displacement. Keeping it as a tensor input lets the - # compiled graph stay reusable across frames. _saved = _patch_task_bufs(task_buf_vals) try: return self.core_compute( @@ -2146,7 +2032,6 @@ def compute_fn( # type: ignore[misc] fparam=fp, aparam=ap, charge_spin=charge_spin, - extended_coord_corr=extended_coord_corr, embedding_only=embedding_only, ) finally: @@ -2222,10 +2107,6 @@ def compute_fn( # type: ignore[misc] spin_for_trace = trace_pad_dim(spin[:1], 0, trace_nf) spin_for_trace = trace_pad_dim(spin_for_trace, 1, trace_nloc) trace_args.append(spin_for_trace) - elif extended_coord_corr is not None: - corr_for_trace = trace_pad_dim(extended_coord_corr[:1], 0, trace_nf) - corr_for_trace = trace_pad_dim(corr_for_trace, 1, trace_nscatter) - trace_args.append(corr_for_trace) # Append task-buffer values last so they map to the *task_buf_vals # varargs in compute_fn. Their shapes are static (they don't vary # batch-to-batch), so passing the actual tensors is correct; make_fx @@ -2306,7 +2187,7 @@ def compute_fn( # type: ignore[misc] # from parameter discovery (FSDP2/DDP would otherwise shard or # synchronise the wrapper's duplicated flat parameter views and # silently corrupt training). The cache is keyed on - # ``(training, has_coord_corr)`` so that distinct graph topologies + # ``training`` so that distinct graph topologies # coexist without evicting each other on every ``model.eval()`` / # ``model.train()`` switch. # NOTE: ``dynamic=True`` emits a single kernel per traced @@ -2764,6 +2645,19 @@ def use_self_built_nlist(self) -> bool: """ return True + def set_inference_neighbor_list(self, neighbor_list: NeighborList | None) -> None: + """Select inference geometry construction while retaining native forward. + + The strategy is runtime state, independent of model parameters and + serialization. Training retains its own batch-aware builder policy. + """ + self._inference_neighbor_list = neighbor_list + + def _get_neighbor_builder(self) -> NeighborList: + if not self.training and self._inference_neighbor_list is not None: + return self._inference_neighbor_list + return AdaptiveNeighborList(method="native") + def build_neighbor_list( self, coord: Float[Tensor, "nf nloc 3"] | Float[Tensor, "nf nloc_x3"], @@ -2772,7 +2666,7 @@ def build_neighbor_list( ) -> EdgeNeighborList: """Build the unified edge-vector schema for the ``forward`` entry.""" nf, nloc = atype.shape[:2] - return _select_neighbor_builder(nf, coord.device).build( + return self._get_neighbor_builder().build( coord.view(nf, nloc, 3), atype, box, @@ -2788,18 +2682,18 @@ def build_extended_neighbor_list( box: Float[Tensor, "nf 9"] | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ - Build the historical extended-coordinate representation. + Build the extended-coordinate representation. This helper is retained for modes whose physical preprocessing still - requires explicit extended atoms, such as the denoising path and the spin - virtual-atom transform. The conservative SeZM energy path converts the + requires explicit extended atoms, such as the denoising path. + The conservative SeZM energy path converts the result to :class:`EdgeNeighborList` before entering ``core_compute``. Returns the :class:`~deepmd.dpmodel.utils.neighbor_list.NeighborList` contract order ``(extended_coord, extended_atype, nlist, mapping)``. """ nf, nloc = atype.shape[:2] - return _select_neighbor_builder(nf, coord.device).build( + return self._get_neighbor_builder().build( coord.view(nf, nloc, 3), atype, box, @@ -3165,9 +3059,7 @@ def supports_edge_parallel(self) -> bool: are completed across ranks (``_gate_partial_exchange``) and the analytical ZBL injection reads extended types, so the parallel path reproduces the folded one. The native spin scheme reuses the - edge_vec interface and therefore participates; only the deepspin - (virtual-atom) scheme uses the nlist interface and is excluded by - the freeze entry point's edge_vec gate. + edge_vec interface and therefore participates. """ descriptor = self.atomic_model.descriptor return bool(descriptor.supports_edge_parallel()) @@ -3179,10 +3071,7 @@ def export_lower_input_kind(self) -> str: (``coord``, ``atype``, ``edge_index``, ``edge_vec``, ``edge_scatter_index``, ``edge_mask``) and the neighbor topology is built on the C++ side, matching :class:`DeepPotPTExpt`. The native - spin scheme reuses this ABI with one extra per-local-atom spin input; - only the deepspin (virtual-atom) scheme overrides it to ``"nlist"`` - because it expands virtual atoms inside the graph from the extended - neighbor list. + spin scheme reuses this ABI with one extra per-local-atom spin input. """ return "edge_vec" diff --git a/deepmd/pt/model/model/sezm_native_spin_model.py b/deepmd/pt/model/model/sezm_native_spin_model.py index 9cc14355b9..92b40a8676 100644 --- a/deepmd/pt/model/model/sezm_native_spin_model.py +++ b/deepmd/pt/model/model/sezm_native_spin_model.py @@ -1,12 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Native-spin SeZM energy model. -Unlike the virtual-atom :class:`SeZMSpinModel`, the native scheme injects the +The model injects the per-atom spin vector directly into the descriptor as an equivariant feature (``l = 0`` magnitude and ``l = 1`` direction) and obtains the magnetic force as -the negative spin gradient of the energy. No virtual atoms are created, so the -neighbor list, type map and selection stay at their real-system sizes, and the -analytical bridging potential needs no real/virtual masking. +the negative spin gradient of the energy. The neighbor list, type map and +selection use the real-system sizes. """ from copy import ( @@ -46,7 +45,7 @@ @BaseModel.register("sezm_native_spin") class SeZMNativeSpinModel(SeZMModel): - """SeZM energy model with native (virtual-atom-free) spin. + """SeZM energy model with native spin. The per-atom spin enters the descriptor through :class:`SpinEmbedding`, and the magnetic force is the negative spin gradient of the energy, produced by @@ -421,8 +420,7 @@ def _attach_spin_masks( enter the descriptor); padding the ghost slots with zero lets it share the extended reduce / fold-back contract that ``communicate_extended_output`` and the LAMMPS C++ backend apply to the - conservative force, so the native and deepspin schemes share one - downstream path. The padding is unconditional (``nall - nloc`` is zero + conservative force. The padding is unconditional (``nall - nloc`` is zero for an isolated cluster) so the closure stays free of shape-dependent branches under ``make_fx`` symbolic tracing. """ diff --git a/deepmd/pt/model/model/sezm_property_model.py b/deepmd/pt/model/model/sezm_property_model.py index e41cfddc78..551b7b71b4 100644 --- a/deepmd/pt/model/model/sezm_property_model.py +++ b/deepmd/pt/model/model/sezm_property_model.py @@ -131,7 +131,6 @@ def core_compute( charge_spin: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, extended_atype: torch.Tensor | None = None, - extended_coord_corr: torch.Tensor | None = None, spin: torch.Tensor | None = None, embedding_only: bool = False, atomic_output_only: bool = False, @@ -149,7 +148,6 @@ def core_compute( charge_spin=charge_spin, comm_dict=comm_dict, extended_atype=extended_atype, - extended_coord_corr=extended_coord_corr, spin=spin, embedding_only=embedding_only, conservative=False, diff --git a/deepmd/pt/model/model/sezm_spin_model.py b/deepmd/pt/model/model/sezm_spin_model.py deleted file mode 100644 index 86b7eeca3d..0000000000 --- a/deepmd/pt/model/model/sezm_spin_model.py +++ /dev/null @@ -1,700 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -"""Spin-enabled SeZM energy model.""" - -import functools -from collections.abc import ( - Callable, -) -from copy import ( - deepcopy, -) -from typing import ( - Any, -) - -import torch - -from deepmd.dpmodel import ( - ModelOutputDef, -) -from deepmd.pt.model.atomic_model.sezm_atomic_model import ( - SeZMAtomicModel, -) -from deepmd.pt.model.descriptor.sezm_nn import ( - nvtx_range, -) -from deepmd.pt.model.model.model import ( - BaseModel, -) -from deepmd.pt.model.model.sezm_model import ( - SeZMModel, -) -from deepmd.pt.model.model.spin_model import ( - SpinModel, - _lookup_type_values, - _pack_spin_stat_sample, -) -from deepmd.pt.utils.utils import ( - to_torch_tensor, -) -from deepmd.pt_expt.utils.edge_schema import ( - edge_schema_from_extended, -) -from deepmd.utils.path import ( - DPPath, -) -from deepmd.utils.spin import ( - Spin, -) -from deepmd.utils.version import ( - check_version_compatibility, -) - - -@BaseModel.register("sezm_spin") -class SeZMSpinModel(SeZMModel): - """SeZM energy model with virtual spin atoms. - - Parameters - ---------- - spin - Spin metadata describing magnetic real types and virtual displacement - scales. - *args - Positional arguments forwarded to :class:`SeZMModel`. - **kwargs - Keyword arguments forwarded to :class:`SeZMModel`. - """ - - model_type = "sezm_spin" - - def __init__( - self, - *args: Any, - spin: Spin, - real_sel: list[int], - **kwargs: Any, - ) -> None: - # A magnetic atom lies inside the bridging window of its own virtual - # partner, so the window mechanics would freeze every magnetic atom; - # the native spin scheme carries no virtual atoms and is the bridged - # spin model. - if str(kwargs.get("bridging_method", "none")).upper() != "NONE": - raise ValueError( - "Analytical bridging is not supported by the deepspin " - "(virtual-atom) scheme; use spin scheme 'native' for a bridged " - "DPA4/SeZM spin model." - ) - super().__init__(*args, **kwargs) - self.spin = spin - self.ntypes_real = self.spin.ntypes_real - self.real_sel = [int(sel) for sel in real_sel] - self.register_buffer( - "virtual_scale_mask", - to_torch_tensor(self.spin.get_virtual_scale_mask()), - persistent=False, - ) - self.register_buffer( - "spin_mask", - to_torch_tensor(self.spin.get_spin_mask()), - persistent=False, - ) - - # ========================================================================= - # Forward Methods - # ========================================================================= - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - spin: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - charge_spin: torch.Tensor | None = None, - ) -> dict[str, torch.Tensor]: - """Return spin-aware SeZM predictions with public output keys.""" - model_ret = self.forward_common( - coord, - atype, - spin, - box=box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - charge_spin=charge_spin, - ) - model_predict: dict[str, torch.Tensor] = { - "atom_energy": model_ret["energy"], - "energy": model_ret["energy_redu"], - "mask_mag": model_ret["mask_mag"], - } - if self.do_grad_r("energy"): - model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) - model_predict["force_mag"] = model_ret["energy_derv_r_mag"].squeeze(-2) - if self.do_grad_c("energy"): - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial: - model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) - return model_predict - - def forward_common( - self, - coord: torch.Tensor, - atype: torch.Tensor, - spin: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - charge_spin: torch.Tensor | None = None, - atomic_output_only: bool = False, - ) -> dict[str, torch.Tensor]: - """Return spin-aware SeZM predictions with internal output keys.""" - with nvtx_range("SeZMSpin/forward_common"): - cc, bb, fp, ap, input_prec = self._input_type_cast( - coord, box=box, fparam=fparam, aparam=aparam - ) - del coord, box, fparam, aparam - atype = atype.to(device=cc.device, dtype=torch.long) - nf, nloc = atype.shape[:2] - if cc.ndim == 2: - cc = cc.view(nf, nloc, 3) - spin = spin.to(dtype=cc.dtype, device=cc.device).reshape(nf, nloc, 3) - - extended_coord, extended_atype, nlist, mapping = ( - self.build_extended_neighbor_list(cc, atype, bb) - ) - extended_spin = torch.gather( - spin, - 1, - mapping.unsqueeze(-1).expand(-1, -1, 3), - ) - ( - extended_coord_updated, - extended_atype_updated, - nlist_updated, - mapping_updated, - extended_coord_corr, - ) = self.process_spin_input_lower( - extended_coord, - extended_atype, - extended_spin, - nlist, - mapping=mapping, - ) - if ap is not None: - ap = self.expand_aparam(ap, nloc * 2) - nlist_updated = self.format_nlist( - extended_coord_updated, - extended_atype_updated, - nlist_updated, - extra_nlist_sort=self.need_sorted_nlist_for_lower(), - ) - edge_schema = edge_schema_from_extended( - extended_coord_updated, - extended_atype_updated[:, : nloc * 2], - nlist_updated, - mapping_updated, - scatter_to_local=True, - ) - model_ret = super().forward_common_lower( - edge_schema.coord, - edge_schema.atype, - edge_schema.edge_index, - edge_schema.edge_vec, - edge_schema.edge_scatter_index, - edge_schema.edge_mask, - fparam=fp, - aparam=ap, - # Slicing the doubled-extended correction down to the local - # region yields a stride-(2*nall*3, ...) view; the compiled - # virial matmul reshapes it to (N, 1, 3), which ``torch.compile`` - # lowers to ``aten.view`` and rejects on the non-contiguous - # layout. Materialize a contiguous copy, mirroring the local - # coord/atype slices in ``edge_schema_from_extended``. - extended_coord_corr=extended_coord_corr[:, : nloc * 2, :].contiguous(), - charge_spin=charge_spin, - input_prec=input_prec, - atomic_output_only=atomic_output_only, - ) - if atomic_output_only: - return model_ret - return self._split_spin_common_output(model_ret, atype, nloc) - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - extended_spin: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - comm_dict: dict[str, torch.Tensor] | None = None, - charge_spin: torch.Tensor | None = None, - ) -> dict[str, torch.Tensor]: - """Return spin-aware SeZM lower-interface predictions.""" - model_ret = self.forward_common_lower( - extended_coord, - extended_atype, - extended_spin, - nlist, - mapping=mapping, - fparam=fparam, - aparam=aparam, - comm_dict=comm_dict, - charge_spin=charge_spin, - extra_nlist_sort=self.need_sorted_nlist_for_lower(), - ) - model_predict: dict[str, torch.Tensor] = { - "atom_energy": model_ret["energy"], - "energy": model_ret["energy_redu"], - "extended_mask_mag": model_ret["mask_mag"], - } - if self.do_grad_r("energy"): - model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) - model_predict["extended_force_mag"] = model_ret[ - "energy_derv_r_mag" - ].squeeze(-2) - if self.do_grad_c("energy"): - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial: - model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( - -2 - ) - return model_predict - - def forward_common_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - extended_spin: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - comm_dict: dict[str, torch.Tensor] | None = None, - extra_nlist_sort: bool = False, - charge_spin: torch.Tensor | None = None, - ) -> dict[str, torch.Tensor]: - """Return spin-aware lower-interface predictions with internal keys.""" - extended_atype = extended_atype.to( - device=extended_coord.device, dtype=torch.long - ) - _, nloc = nlist.shape[:2] - ( - extended_coord_updated, - extended_atype_updated, - nlist_updated, - mapping_updated, - extended_coord_corr, - ) = self.process_spin_input_lower( - extended_coord, - extended_atype, - extended_spin, - nlist, - mapping=mapping, - ) - if aparam is not None: - aparam = self.expand_aparam(aparam, nloc * 2) - nlist_updated = self.format_nlist( - extended_coord_updated, - extended_atype_updated, - nlist_updated, - extra_nlist_sort=extra_nlist_sort, - ) - edge_schema = edge_schema_from_extended( - extended_coord_updated, - extended_atype_updated[:, : nloc * 2], - nlist_updated, - mapping_updated, - ) - model_ret = super().forward_common_lower( - edge_schema.coord, - edge_schema.atype, - edge_schema.edge_index, - edge_schema.edge_vec, - edge_schema.edge_scatter_index, - edge_schema.edge_mask, - fparam=fparam, - aparam=aparam, - comm_dict=comm_dict, - extended_coord_corr=extended_coord_corr, - charge_spin=charge_spin, - ) - return self._split_spin_lower_output(model_ret, extended_atype, nloc) - - def forward_common_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - extended_spin: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - ) -> torch.nn.Module: - """Trace the spin lower interface into an exportable FX graph.""" - extra_sort = self.need_sorted_nlist_for_lower() - - def lower_fn( - ext_coord: torch.Tensor, - ext_atype: torch.Tensor, - ext_spin: torch.Tensor, - nlist_: torch.Tensor, - mapping_: torch.Tensor | None, - fparam_: torch.Tensor | None, - aparam_: torch.Tensor | None, - charge_spin_: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - ext_coord = ext_coord.detach().requires_grad_(True) - return self.forward_common_lower( - ext_coord, - ext_atype, - ext_spin, - nlist_, - mapping_, - fparam=fparam_, - aparam=aparam_, - extra_nlist_sort=extra_sort, - charge_spin=charge_spin_, - ) - - def fn( - ext_coord: torch.Tensor, - ext_atype: torch.Tensor, - ext_spin: torch.Tensor, - nlist_: torch.Tensor, - mapping_: torch.Tensor | None, - fparam_: torch.Tensor | None, - aparam_: torch.Tensor | None, - charge_spin_: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - return lower_fn( - ext_coord, - ext_atype, - ext_spin, - nlist_, - mapping_, - fparam_, - aparam_, - charge_spin_, - ) - - trace_inputs = ( - extended_coord, - extended_atype, - extended_spin, - nlist, - mapping, - fparam, - aparam, - ) - if self.get_dim_chg_spin() > 0: - charge_spin = self.convert_charge_spin( - charge_spin, - nf=extended_atype.shape[0], - dtype=extended_coord.dtype, - device=extended_coord.device, - ) - trace_inputs = (*trace_inputs, charge_spin) - - return self.trace_lower_exportable( - fn, - *trace_inputs, - ) - - # ========================================================================= - # Statistics and Mode Methods - # ========================================================================= - - def compute_or_load_stat( - self, - sampled_func: Callable[[], list[dict[str, Any]]], - stat_file_path: DPPath | None = None, - preset_observed_type: list[str] | None = None, - ) -> None: - """Compute or load statistics with virtual spin atoms included.""" - super().compute_or_load_stat( - self._get_spin_sampled_func(sampled_func), - stat_file_path, - preset_observed_type=preset_observed_type, - ) - - def change_out_bias( - self, - merged: Callable[[], list[dict[str, Any]]] | list[dict[str, Any]], - bias_adjust_mode: str = "change-by-statistic", - ) -> None: - """Change output bias using spin-expanded sampled data.""" - spin_sampled_func = self._get_spin_sampled_func( - merged if callable(merged) else lambda: merged - ) - super().change_out_bias( - spin_sampled_func, - bias_adjust_mode=bias_adjust_mode, - ) - - def change_type_map( - self, type_map: list[str], model_with_new_type_stat: Any = None - ) -> None: - """Change real type map and rebuild corresponding virtual spin types.""" - type_map_with_spin = type_map + [item + "_spin" for item in type_map] - super().change_type_map(type_map_with_spin, model_with_new_type_stat) - self.ntypes_real = len(type_map) - - def set_active_mode(self, mode: str) -> None: - """Switch mode, allowing only the conservative energy path.""" - normalized = str(mode).lower() - if normalized != "ener": - raise NotImplementedError("SeZM spin supports only the `ener` path.") - super().set_active_mode(normalized) - - def set_active_mode_from_loss(self, loss_type: str) -> None: - """Select execution mode from loss type.""" - normalized = str(loss_type).lower() - if normalized == "dens": - raise NotImplementedError("SeZM spin supports only the `ener` path.") - if normalized in {"ener", "ener_spin"}: - self.set_active_mode("ener") - - # ========================================================================= - # Output Definitions and Metadata - # ========================================================================= - - def has_spin(self) -> bool: - """Return whether this model consumes spin input.""" - return True - - def export_lower_input_kind(self) -> str: - """Return the ``.pt2`` lower ABI: the deepspin scheme needs the nlist. - - Virtual atoms are placed and the neighbor list is expanded inside the - traced graph from the extended inputs, so the export contract feeds the - real extended coordinates, spin and neighbor list rather than a - pre-built edge schema. - """ - return "nlist" - - def get_type_map(self) -> list[str]: - """Return the real atom type map.""" - return super().get_type_map()[: self.ntypes_real] - - def get_ntypes(self) -> int: - """Return the number of real atom types.""" - return len(self.get_type_map()) - - def get_sel(self) -> list[int]: - """Return the public real-atom neighbor selection.""" - return self.real_sel - - def get_nsel(self) -> int: - """Return the public real-atom total neighbor count.""" - return int(sum(self.real_sel)) - - def get_nnei(self) -> int: - """Return the public real-atom total neighbor count.""" - return int(sum(self.real_sel)) - - def get_observed_type_list(self) -> list[str]: - """Return observed real types according to the output bias.""" - type_map = self.get_type_map() - out_bias = self.atomic_model.get_out_bias()[0] - assert out_bias is not None, "No out_bias found in the model." - assert out_bias.dim() == 2, "The supported out_bias should be a 2D tensor." - assert out_bias.size(0) >= self.ntypes_real, ( - "The out_bias shape is smaller than the number of real types." - ) - bias_mask = ( - torch.gt(torch.abs(out_bias[: self.ntypes_real]), 1e-6).any(dim=-1).cpu() - ) - result: list[str] = [] - for t, m in zip(type_map, bias_mask.tolist()): - if m: - result.append(t) - return result - - def model_output_def(self) -> ModelOutputDef: - """Return the spin-aware model output definition.""" - atomic_output_def = self.atomic_output_def() - atomic_output_def["energy"].magnetic = True - return ModelOutputDef(atomic_output_def) - - def translated_output_def(self) -> dict[str, Any]: - """Translate internal output definitions to public spin keys.""" - out_def_data = self.model_output_def().get_data() - output_def = { - "atom_energy": out_def_data["energy"], - "energy": out_def_data["energy_redu"], - "mask_mag": out_def_data["mask_mag"], - } - if self.do_grad_r("energy"): - output_def["force"] = deepcopy(out_def_data["energy_derv_r"]) - output_def["force"].squeeze(-2) - output_def["force_mag"] = deepcopy(out_def_data["energy_derv_r_mag"]) - output_def["force_mag"].squeeze(-2) - if self.do_grad_c("energy"): - output_def["virial"] = deepcopy(out_def_data["energy_derv_c_redu"]) - output_def["virial"].squeeze(-2) - output_def["atom_virial"] = deepcopy(out_def_data["energy_derv_c"]) - output_def["atom_virial"].squeeze(-2) - return output_def - - # ========================================================================= - # Serialization - # ========================================================================= - - def serialize(self) -> dict[str, Any]: - """Serialize the SeZM spin model.""" - data = super().serialize() - data["type"] = self.model_type - data["spin"] = self.spin.serialize() - data["real_sel"] = self.real_sel - return data - - @classmethod - def deserialize(cls, data: dict[str, Any]) -> "SeZMSpinModel": - """Deserialize a SeZM spin model.""" - data = data.copy() - version = int(data.pop("@version", 1)) - check_version_compatibility(version, 1, 1) - data.pop("@class", None) - data.pop("type", None) - spin = Spin.deserialize(data.pop("spin")) - real_sel = data.pop("real_sel") - atomic_model = SeZMAtomicModel.deserialize(data.pop("atomic_model")) - return cls(atomic_model_=atomic_model, spin=spin, real_sel=real_sel, **data) - - # ========================================================================= - # Small Utilities - # ========================================================================= - - def build_neighbor_list( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Build the real-atom neighbor list before spin expansion.""" - return super().build_neighbor_list(coord, atype, box) - - def format_nlist( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - extra_nlist_sort: bool = False, - ) -> torch.Tensor: - """Format spin-expanded nlist to the internal descriptor capacity.""" - del extended_atype - return self._format_nlist( - extended_coord, - nlist, - sum(self.atomic_model.get_sel()), - extra_nlist_sort=extra_nlist_sort, - ) - - def _get_spin_sampled_func( - self, sampled_func: Callable[[], list[dict[str, Any]]] - ) -> Callable[[], list[dict[str, Any]]]: - """Wrap a data sampler so statistics see real and virtual atoms.""" - - @functools.lru_cache - def spin_sampled_func() -> list[dict[str, Any]]: - packed_samples = [] - for sample in sampled_func(): - packed = _pack_spin_stat_sample(self, sample) - packed["model_coord"] = sample["coord"] - packed["model_atype"] = sample["atype"] - packed["model_spin"] = sample["spin"] - if "aparam" in sample: - packed["model_aparam"] = sample["aparam"] - packed_samples.append(packed) - return packed_samples - - return self.atomic_model._make_wrapped_sampler(spin_sampled_func) - - def _ensure_mask_mag( - self, - model_ret: dict[str, torch.Tensor], - atype: torch.Tensor, - ) -> None: - """Ensure the magnetic atom mask exists in ``model_ret``.""" - if "mask_mag" in model_ret: - return - nframes, nloc = atype.shape[:2] - atomic_mask = _lookup_type_values(self.virtual_scale_mask, atype).reshape( - [nframes, nloc, 1] - ) - model_ret["mask_mag"] = atomic_mask > 0.0 - - def _split_spin_common_output( - self, - model_ret: dict[str, torch.Tensor], - atype: torch.Tensor, - nloc: int, - ) -> dict[str, torch.Tensor]: - """Split full-interface SeZM outputs into real and magnetic parts.""" - model_ret["energy"] = torch.split(model_ret["energy"], [nloc, nloc], dim=1)[0] - if self.do_grad_r("energy") and model_ret.get("energy_derv_r") is not None: - ( - model_ret["energy_derv_r"], - model_ret["energy_derv_r_mag"], - model_ret["mask_mag"], - ) = self.process_spin_output(atype, model_ret["energy_derv_r"]) - if self.do_grad_c("energy") and model_ret.get("energy_derv_c") is not None: - ( - model_ret["energy_derv_c"], - model_ret["energy_derv_c_mag"], - model_ret["mask_mag"], - ) = self.process_spin_output( - atype, - model_ret["energy_derv_c"], - add_mag=True, - virtual_scale=False, - ) - self._ensure_mask_mag(model_ret, atype) - return model_ret - - def _split_spin_lower_output( - self, - model_ret: dict[str, torch.Tensor], - extended_atype: torch.Tensor, - nloc: int, - ) -> dict[str, torch.Tensor]: - """Split lower-interface SeZM outputs into real and magnetic parts.""" - model_ret["energy"] = torch.split(model_ret["energy"], [nloc, nloc], dim=1)[0] - if self.do_grad_r("energy") and model_ret.get("energy_derv_r") is not None: - ( - model_ret["energy_derv_r"], - model_ret["energy_derv_r_mag"], - model_ret["mask_mag"], - ) = self.process_spin_output_lower( - extended_atype, model_ret["energy_derv_r"], nloc - ) - if self.do_grad_c("energy") and model_ret.get("energy_derv_c") is not None: - ( - model_ret["energy_derv_c"], - model_ret["energy_derv_c_mag"], - model_ret["mask_mag"], - ) = self.process_spin_output_lower( - extended_atype, - model_ret["energy_derv_c"], - nloc, - add_mag=True, - virtual_scale=False, - ) - self._ensure_mask_mag(model_ret, extended_atype) - return model_ret - - process_spin_input = SpinModel.process_spin_input - process_spin_input_lower = SpinModel.process_spin_input_lower - process_spin_output = SpinModel.process_spin_output - process_spin_output_lower = SpinModel.process_spin_output_lower - extend_nlist = staticmethod(SpinModel.extend_nlist) - expand_aparam = staticmethod(SpinModel.expand_aparam) diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index 550de0bfa4..fe8c0434e0 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -18,6 +18,9 @@ from deepmd.pt.model.atomic_model import ( DPAtomicModel, ) +from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, +) from deepmd.pt.utils.spin import ( concat_switch_virtual, ) @@ -27,6 +30,10 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.preset_out_bias import ( + normalize_preset_out_bias, + remap_preset_out_bias, +) from deepmd.utils.spin import ( Spin, ) @@ -124,6 +131,9 @@ def __init__( backbone_model: DPAtomicModel, spin: Spin, ) -> None: + """Initialize a spin wrapper around a compatible backbone.""" + if any(isinstance(module, DescrptSeZM) for module in backbone_model.modules()): + raise ValueError("DPA4/SeZM spin models require a native spin model.") super().__init__() self.backbone_model = backbone_model self.spin = spin @@ -518,6 +528,7 @@ def change_out_bias( self, merged: Callable[[], list[dict]] | list[dict], bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. @@ -535,13 +546,24 @@ def change_out_bias( 'change-by-statistic' : perform predictions on labels of target dataset, and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. + preset_out_bias : dict, optional + Preset on the public real-atom type map. It is normalized and + extended with unassigned virtual types before forwarding. Omit + to retain the model preset; an empty dict fits without fixing + preset outputs. Non-finite stored cells covered by a checkpoint + preset are recovered before residual fitting. """ + preset_out_bias = remap_preset_out_bias( + normalize_preset_out_bias(preset_out_bias, self.get_type_map()), + list(range(self.get_ntypes())) + [-1] * self.get_ntypes(), + ) spin_sampled_func = self._get_spin_sampled_func( merged if callable(merged) else lambda: merged ) self.backbone_model.change_out_bias( spin_sampled_func, bias_adjust_mode=bias_adjust_mode, + preset_out_bias=preset_out_bias, ) def change_type_map( diff --git a/deepmd/pt/model/model/transform_output.py b/deepmd/pt/model/model/transform_output.py index 3a528c1f57..2369d8c724 100644 --- a/deepmd/pt/model/model/transform_output.py +++ b/deepmd/pt/model/model/transform_output.py @@ -219,7 +219,6 @@ def edge_energy_deriv( nf: int, nall: int, create_graph: bool, - extended_coord_corr: torch.Tensor | None = None, spin_leaf: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: """Assemble extended force, virial and atomic virial from edge gradients. @@ -263,9 +262,6 @@ def edge_energy_deriv( create_graph Keep the first-derivative graph alive so the force-loss second backward can reach the parameters. - extended_coord_corr - Optional spin virtual-displacement correction with shape - ``(nf, nall, 3)``; adds ``force (x) coord_corr`` per extended atom. spin_leaf Optional per-atom spin leaf with shape ``(nf, nloc, 3)`` for the native spin scheme. When provided, the energy is also differentiated with @@ -385,17 +381,9 @@ def edge_energy_deriv( av_flat = av_flat.index_add(0, src_ext, w_edge) extended_virial = av_flat.view(nf, nall, 9) - if extended_coord_corr is not None: - # Spin: the virtual-atom displacement adds force (x) coord_corr per atom. - corr = ( - extended_force.unsqueeze(-1) - @ extended_coord_corr.unsqueeze(-2).to(extended_force.dtype) - ).reshape(nf, nall, 9) - extended_virial = extended_virial + corr - energy_derv_r = extended_force.unsqueeze(-2) energy_derv_c = extended_virial.unsqueeze(-2) - if frame_virial is not None and extended_coord_corr is None: + if frame_virial is not None: energy_derv_c_redu = frame_virial.to(env.GLOBAL_PT_ENER_FLOAT_PRECISION).view( nf, 1, 9 ) @@ -405,8 +393,7 @@ def edge_energy_deriv( ) # Magnetic force is the negative spin gradient, matching the dataset - # ``force_mag = -dE/dspin`` convention (the virtual-atom scheme reaches the - # same quantity through ``F_virtual * virtual_scale``). + # ``force_mag = -dE/dspin`` convention. energy_derv_r_mag: torch.Tensor | None = None if spin_leaf is not None: energy_derv_r_mag = (-grads[1]).unsqueeze(-2) diff --git a/deepmd/pt/model/task/__init__.py b/deepmd/pt/model/task/__init__.py index f9f8e6561d..c6e7df38ed 100644 --- a/deepmd/pt/model/task/__init__.py +++ b/deepmd/pt/model/task/__init__.py @@ -2,6 +2,9 @@ from .base_fitting import ( BaseFitting, ) +from .density import ( + DensityFittingNet, +) from .dipole import ( DipoleFittingNet, ) @@ -31,6 +34,7 @@ __all__ = [ "BaseFitting", "DOSFittingNet", + "DensityFittingNet", "DipoleFittingNet", "EnergyFittingNet", "EnergyFittingNetDirect", diff --git a/deepmd/pt/model/task/density.py b/deepmd/pt/model/task/density.py new file mode 100644 index 0000000000..5f386ef540 --- /dev/null +++ b/deepmd/pt/model/task/density.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.pt.model.task.fitting import ( + Fitting, + GeneralFitting, +) +from deepmd.pt.model.task.invar_fitting import ( + InvarFitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE + +log = logging.getLogger(__name__) + + +@Fitting.register("density") +class DensityFittingNet(InvarFitting): + def __init__( + self, + ntypes: int, + dim_descrpt: int, + neuron: list[int] = [128, 128, 128], + bias_atom_e: torch.Tensor | None = None, + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + mixed_types: bool = True, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + **kwargs: Any, + ) -> None: + super().__init__( + "density", + ntypes, + dim_descrpt, + 1, + neuron=neuron, + bias_atom_e=bias_atom_e, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + activation_function=activation_function, + precision=precision, + mixed_types=mixed_types, + seed=seed, + type_map=type_map, + **kwargs, + ) + if numb_aparam > 0: + raise ValueError( + "density fitting does not support atomic parameters (aparam): " + "the fitting net consumes the grid-point descriptor rows, " + "which have no per-atom parameters" + ) + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + @classmethod + def deserialize(cls, data: dict) -> "GeneralFitting": + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 5, 1) + data.pop("var_name") + data.pop("dim_out") + return super().deserialize(data) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + **super().serialize(), + "type": "density", + } + + # make jit happy with torch 2.0.0 + exclude_types: list[int] diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 7efc76e4d9..1f1f4c6029 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -35,8 +35,8 @@ change_model_out_bias, resolve_step_schedule, ) -from deepmd.dpmodel.utils import ( - compute_total_numb_batch, +from deepmd.dpmodel.train.data import ( + add_training_data_requirements, ) from deepmd.loggers import ( is_node_main_process, @@ -51,6 +51,7 @@ DOSLoss, EnergySpinLoss, EnergyStdLoss, + GridDensityLoss, PopulationLoss, PropertyLoss, TaskLoss, @@ -85,21 +86,20 @@ ) from deepmd.pt.utils.dataloader import ( DpLoaderSet, - get_sampler_from_params, + create_dataloader, ) from deepmd.pt.utils.env import ( DEVICE, JIT, LOCAL_RANK, - NUM_WORKERS, SAMPLER_RECORD, ) from deepmd.pt.utils.learning_rate import ( BaseLR, ) -from deepmd.pt.utils.lmdb_dataset import ( - LmdbBatchDataLoader, - LmdbDataset, +from deepmd.pt.utils.spin_finetune import ( + compute_native_spin_input_stats, + get_native_spin_finetune_state, ) from deepmd.pt.utils.stat import ( make_stat_input, @@ -107,9 +107,6 @@ scan_redu_stats, select_batch_frames, ) -from deepmd.pt.utils.utils import ( - to_numpy_array, -) from deepmd.pt_expt.train.ema import ( EMA_CHECKPOINT_KEY, ModelEMA, @@ -132,7 +129,11 @@ from deepmd.utils.data import ( DataRequirementItem, has_data_requirement, - min_pair_dist_requirement, +) +from deepmd.utils.data_group import ( + GroupedTrainingDataLoader, + TrainingDataLoader, + iter_data_sources, ) from deepmd.utils.finetune import ( warn_configuration_mismatch_during_finetune, @@ -162,9 +163,6 @@ ZeroRedundancyOptimizer, ) from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.data import ( - DataLoader, -) from deepmd.utils.stat_file import ( StatFileSpec, @@ -289,103 +287,36 @@ def cycle_iterator(iterable: Iterable) -> Generator[Any, None, None]: yield from it def get_data_loader( - _training_data: DpLoaderSet | LmdbDataset, - _validation_data: DpLoaderSet | LmdbDataset | None, + _training_data: Any, + _validation_data: Any | None, _training_params: dict[str, Any], ) -> tuple[ - DataLoader | LmdbBatchDataLoader, + TrainingDataLoader | GroupedTrainingDataLoader, Generator[Any, None, None], - DataLoader | LmdbBatchDataLoader | None, + TrainingDataLoader | GroupedTrainingDataLoader | None, Generator[Any, None, None] | None, int, ]: - def get_dataloader_and_iter_lmdb( - _data: LmdbDataset, - ) -> tuple[LmdbBatchDataLoader, Generator[Any, None, None]]: - # The sampler yields one list[int] per batch. Every batch is - # homogeneous in label availability; whether it is also - # homogeneous in atom count follows from the dataset's - # batch_size rule, which the sampler reads off the reader. - from deepmd.dpmodel.utils.lmdb_data import ( - LmdbBatchSampler, - ) - - _block_targets = getattr(_data, "_block_targets", None) - _sampler_kwargs = { - "shuffle": True, - "seed": _training_params.get("seed"), - "block_targets": _block_targets, - } - - if self.world_size > 1: - from deepmd.dpmodel.utils.lmdb_data import ( - DistributedLmdbBatchSampler, - ) - - _inner_sampler = DistributedLmdbBatchSampler( - _data._reader, - rank=self.rank, - world_size=self.world_size, - **_sampler_kwargs, - ) - else: - _inner_sampler = LmdbBatchSampler( - _data._reader, - **_sampler_kwargs, - ) - - _dataloader = LmdbBatchDataLoader( - _data, - _inner_sampler, - pin_memory=DEVICE.type != "cpu", - ) - _data_iter = cycle_iterator(_dataloader) - return _dataloader, _data_iter - def get_dataloader_and_iter( - _data: DpLoaderSet, _params: dict[str, Any] - ) -> tuple[DataLoader, Generator[Any, None, None]]: - _sampler = get_sampler_from_params(_data, _params) - if _sampler is None: - log.warning( - "Sampler not specified!" - ) # None sampler will lead to a premature stop iteration. Replacement should be True in attribute of the sampler to produce expected number of items in one iteration. - _dataloader = DataLoader( + _data: Any, _params: Any + ) -> tuple[Any, Generator[Any, None, None]]: + _dataloader = create_dataloader( _data, - sampler=_sampler, - batch_size=None, - num_workers=NUM_WORKERS - if dist.is_available() - else 0, # setting to 0 diverges the behavior of its iterator; should be >=1 - drop_last=False, - collate_fn=lambda batch: batch, # prevent extra conversion - pin_memory=(DEVICE != "cpu"), # pin memory only if not on CPU + _params, + seed=_training_params.get("seed"), + rank=self.rank, + world_size=self.world_size, ) - _data_iter = cycle_iterator(_dataloader) - return _dataloader, _data_iter + return _dataloader, cycle_iterator(_dataloader) - if isinstance(_training_data, LmdbDataset): - training_dataloader, training_data_iter = get_dataloader_and_iter_lmdb( - _training_data - ) - else: - training_dataloader, training_data_iter = get_dataloader_and_iter( - _training_data, _training_params["training_data"] - ) + training_dataloader, training_data_iter = get_dataloader_and_iter( + _training_data, _training_params["training_data"] + ) if _validation_data is not None: - if isinstance(_validation_data, LmdbDataset): - ( - validation_dataloader, - validation_data_iter, - ) = get_dataloader_and_iter_lmdb(_validation_data) - else: - ( - validation_dataloader, - validation_data_iter, - ) = get_dataloader_and_iter( - _validation_data, _training_params["validation_data"] - ) + validation_dataloader, validation_data_iter = get_dataloader_and_iter( + _validation_data, _training_params["validation_data"] + ) valid_numb_batch = _training_params["validation_data"].get( "numb_btch", 1 ) @@ -404,9 +335,9 @@ def get_dataloader_and_iter( def single_model_stat( _model: Any, _data_stat_nbatch: int, - _training_data: DpLoaderSet, + _training_data: Any, _stat_file_spec: StatFileSpec, - _min_pair_dist: float = 0.0, + _filter_enabled: bool = False, _data_stat_full: bool = False, finetune_has_new_type: bool = False, preset_observed_type: list[str] | None = None, @@ -417,7 +348,10 @@ def get_sample() -> Any: _training_data.systems, _training_data.dataloaders, _data_stat_nbatch, - min_pair_dist=_min_pair_dist, + min_pair_dist=_filter_enabled, + filter_enabled_by_system=getattr( + _training_data, "stat_filter_enabled", None + ), ) return sampled @@ -430,7 +364,10 @@ def get_sample() -> Any: ntypes, keys, intensive=intensive, - min_pair_dist=_min_pair_dist, + min_pair_dist=_filter_enabled, + filter_enabled_by_system=getattr( + _training_data, "stat_filter_enabled", None + ), ) ) @@ -520,38 +457,37 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: ) # Data - # Radius of the pair-clearance filter, per task. - self.min_pair_dist_by_task: dict[str, float] = {} + # Per-group requirements own the filter radii; tasks only enable filtering. + self.pair_filter_enabled_by_task: dict[str, bool] = {} if not self.multi_task: # add data requirement for labels data_requirement = self.loss.label_requirement data_requirement += get_additional_data_requirement(self.model) - pair_filter = min_pair_dist_requirement( + filter_enabled = add_training_data_requirements( + training_data, + training_params["training_data"], model_params, - float( - training_params.get("training_data", {}).get("min_pair_dist", 0.0) - ), + data_requirement, + lambda data, requirements: data.add_data_requirement(requirements), ) - min_pair_dist = 0.0 if pair_filter is None else float(pair_filter.default) - self.min_pair_dist_by_task["Default"] = min_pair_dist - if pair_filter is not None: - data_requirement.append(pair_filter) - training_data.add_data_requirement(data_requirement) + self.pair_filter_enabled_by_task["Default"] = filter_enabled if validation_data is not None: validation_data.add_data_requirement( self.loss.label_requirement + get_additional_data_requirement(self.model) ) # Preload and apply modifiers to all data before computing statistics - training_data.preload_and_modify_all_data_torch() + for source in iter_data_sources(training_data): + source.preload_and_modify_all_data_torch() if validation_data is not None: - validation_data.preload_and_modify_all_data_torch() + for source in iter_data_sources(validation_data): + source.preload_and_modify_all_data_torch() self.get_sample_func = single_model_stat( self.model, model_params.get("data_stat_nbatch", 10), training_data, self.stat_file_specs["Default"], - _min_pair_dist=min_pair_dist, + _filter_enabled=filter_enabled, _data_stat_full=model_params.get("data_stat_full", False), finetune_has_new_type=self.finetune_links["Default"].get_has_new_type() if self.finetune_links is not None @@ -571,19 +507,9 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: self.validation_data, self.valid_numb_batch, ) = get_data_loader(training_data, validation_data, training_params) - training_data.print_summary( - "training", - to_numpy_array(self.training_dataloader.sampler.weights) - if not isinstance(training_data, LmdbDataset) - else [1.0], - ) + self.training_dataloader.print_summary("training") if validation_data is not None: - validation_data.print_summary( - "validation", - to_numpy_array(self.validation_dataloader.sampler.weights) - if not isinstance(validation_data, LmdbDataset) - else [1.0], - ) + self.validation_dataloader.print_summary("validation") else: ( self.training_dataloader, @@ -599,30 +525,25 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: data_requirement += get_additional_data_requirement( self.model[model_key] ) - pair_filter = min_pair_dist_requirement( + filter_enabled = add_training_data_requirements( + training_data[model_key], + training_params["data_dict"][model_key]["training_data"], model_params["model_dict"][model_key], - float( - training_params["data_dict"][model_key] - .get("training_data", {}) - .get("min_pair_dist", 0.0) - ), - ) - min_pair_dist = ( - 0.0 if pair_filter is None else float(pair_filter.default) + data_requirement, + lambda data, requirements: data.add_data_requirement(requirements), ) - self.min_pair_dist_by_task[model_key] = min_pair_dist - if pair_filter is not None: - data_requirement.append(pair_filter) - training_data[model_key].add_data_requirement(data_requirement) + self.pair_filter_enabled_by_task[model_key] = filter_enabled if validation_data[model_key] is not None: validation_data[model_key].add_data_requirement( self.loss[model_key].label_requirement + get_additional_data_requirement(self.model[model_key]) ) # Preload and apply modifiers to all data before computing statistics - training_data[model_key].preload_and_modify_all_data_torch() + for source in iter_data_sources(training_data[model_key]): + source.preload_and_modify_all_data_torch() if validation_data[model_key] is not None: - validation_data[model_key].preload_and_modify_all_data_torch() + for source in iter_data_sources(validation_data[model_key]): + source.preload_and_modify_all_data_torch() _mt_user_observed = ( model_params["model_dict"][model_key] .get("info", {}) @@ -633,7 +554,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: model_params["model_dict"][model_key].get("data_stat_nbatch", 10), training_data[model_key], self.stat_file_specs[model_key], - _min_pair_dist=min_pair_dist, + _filter_enabled=filter_enabled, _data_stat_full=model_params["model_dict"][model_key].get( "data_stat_full", False ), @@ -667,40 +588,26 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: training_params["data_dict"][model_key], ) - training_data[model_key].print_summary( - f"training in {model_key}", - to_numpy_array(self.training_dataloader[model_key].sampler.weights) - if not isinstance(training_data[model_key], LmdbDataset) - else [1.0], + self.training_dataloader[model_key].print_summary( + f"training in {model_key}" ) if ( validation_data is not None and validation_data[model_key] is not None ): - validation_data[model_key].print_summary( - f"validation in {model_key}", - to_numpy_array( - self.validation_dataloader[model_key].sampler.weights - ) - if not isinstance(validation_data[model_key], LmdbDataset) - else [1.0], + self.validation_dataloader[model_key].print_summary( + f"validation in {model_key}" ) # Resolve training steps def epoch_length(model_key: str) -> int: """Return the batches this rank consumes in one epoch of a task.""" - _data = training_data[model_key] if self.multi_task else training_data _dataloader = ( self.training_dataloader[model_key] if self.multi_task else self.training_dataloader ) - if isinstance(_data, LmdbDataset): - return len(_dataloader) - return compute_total_numb_batch( - _data.index, - to_numpy_array(_dataloader.sampler.weights), - ) + return _dataloader.epoch_length schedule = resolve_step_schedule( training_params, @@ -733,13 +640,7 @@ def epoch_length(model_key: str) -> int: self.nonfinite_grad_guard = NonFiniteGradGuard() self.lr_schedule = get_lr(config["learning_rate"]) - # Radius the pair-clearance filter was registered with. The - # filter reads it as an enable flag, the frames carrying a margin that - # is already measured against that radius. It is resolved once with the - # data requirement, so the two cannot disagree. - self.has_min_pair_filter = any( - value > 0.0 for value in self.min_pair_dist_by_task.values() - ) + self.has_min_pair_filter = any(self.pair_filter_enabled_by_task.values()) if self.has_min_pair_filter: if self.multi_task: local_training_batch_attempts = max( @@ -874,6 +775,24 @@ def collect_single_finetune_params( _pretrained_model_params["descriptor"], _model_key_from, ) + if _finetune_rule_single.convert_to_native_spin: + log.info( + "Initialize native-spin branch '%s' from spin-free branch '%s'.", + _model_key, + _model_key_from, + ) + converted = get_native_spin_finetune_state( + self.wrapper.model[_model_key], + pretrained_model_wrapper.model[_model_key_from], + random_fitting=_new_fitting, + ) + _new_state_dict.update( + { + f"model.{_model_key}.{key}": value + for key, value in converted.items() + } + ) + return target_keys = [ i for i in _random_state_dict.keys() @@ -950,11 +869,16 @@ def update_finetune_bias( model = ( self.model[_model_key] if self.multi_task else self.model ) - model = model_change_out_bias( - model, + sample_func = ( self.get_sample_func[_model_key] if self.multi_task - else self.get_sample_func, + else self.get_sample_func + ) + if _finetune_rule.convert_to_native_spin: + compute_native_spin_input_stats(model, sample_func) + model = model_change_out_bias( + model, + sample_func, _bias_adjust_mode="change-by-statistic" if not _finetune_rule.get_random_fitting() else "set-by-statistic", @@ -1409,7 +1333,7 @@ def run(self) -> None: try: self._run() finally: - self._close_lmdb_loaders() + self._close_data_loaders() def _training_results( self, more_loss: dict[str, Any], task_key: str = "Default" @@ -1732,6 +1656,7 @@ def log_loss_valid(_task_key: str = "Default") -> dict: step_id=_step_id, display_step=display_step_id, lr=cur_lr, + prepare_checkpoint=self.ensure_finite_gradients_for_checkpoint, save_checkpoint=( self.save_model_merged if self._lora_enabled @@ -1743,6 +1668,7 @@ def log_loss_valid(_task_key: str = "Default") -> dict: step_id=_step_id, display_step=display_step_id, lr=cur_lr, + prepare_checkpoint=self.ensure_finite_gradients_for_checkpoint, save_checkpoint=( self.save_ema_model_merged if self._lora_enabled @@ -1756,27 +1682,9 @@ def log_loss_valid(_task_key: str = "Default") -> dict: and _step_id != self.start_step ) if should_save_checkpoint: - # Abort before writing if any gradient norm since the previous - # checkpoint was non-finite. - self.nonfinite_grad_guard.raise_if_nonfinite( - self.wrapper.named_parameters + self._save_training_checkpoint( + step=_step_id, display_step=display_step_id, lr=cur_lr ) - if should_save_checkpoint and ( - self.sharding.enabled or self.rank == 0 or dist.get_rank() == 0 - ): - # Handle the case if rank 0 aborted and re-assigned - self.latest_model = self.ckpt_store.path_for(display_step_id) - self.save_model(self.latest_model, lr=cur_lr, step=_step_id) - if self.rank == 0 or dist.get_rank() == 0: - log.info(f"Saved model to {self.latest_model}") - self.ckpt_store.publish(self.latest_model) - if self.model_ema is not None: - self.latest_ema_model = self.ema_ckpt_store.path_for( - display_step_id - ) - self.save_ema_model(self.latest_ema_model, lr=cur_lr, step=_step_id) - if self.rank == 0 or dist.get_rank() == 0: - self.ema_ckpt_store.publish(self.latest_ema_model) # tensorboard if self.enable_tensorboard and ( @@ -1844,62 +1752,32 @@ def log_loss_valid(_task_key: str = "Default") -> dict: self._discarded_training_batches, ) - if ( - self.change_bias_after_training - and self.num_steps > self.start_step - and (self.rank == 0 or dist.get_rank() == 0) - ): - if not self.multi_task: - self.model = model_change_out_bias( - self.model, - self.get_sample_func, - _bias_adjust_mode="change-by-statistic", - ) - else: - for model_key in self.model_keys: - self.model[model_key] = model_change_out_bias( - self.model[model_key], - self.get_sample_func[model_key], + if self.change_bias_after_training and self.num_steps > self.start_step: + if self.rank == 0 or dist.get_rank() == 0: + if not self.multi_task: + self.model = model_change_out_bias( + self.model, + self.get_sample_func, _bias_adjust_mode="change-by-statistic", ) - self.latest_model = self.ckpt_store.path_for(self.num_steps) - cur_lr = self.lr_schedule.value(self.num_steps - 1) - self.save_model(self.latest_model, lr=cur_lr, step=self.num_steps - 1) - log.info(f"Saved model to {self.latest_model}") - self.ckpt_store.publish(self.latest_model) - if self.model_ema is not None: - self.latest_ema_model = self.ema_ckpt_store.path_for(self.num_steps) - self.save_ema_model( - self.latest_ema_model, - lr=cur_lr, - step=self.num_steps - 1, - ) - self.ema_ckpt_store.publish(self.latest_ema_model) + else: + for model_key in self.model_keys: + self.model[model_key] = model_change_out_bias( + self.model[model_key], + self.get_sample_func[model_key], + _bias_adjust_mode="change-by-statistic", + ) + self._save_training_checkpoint( + step=self.num_steps - 1, + display_step=self.num_steps, + lr=self.lr_schedule.value(self.num_steps - 1), + ) - if self.num_steps == 0 and self.sharding.enabled: - # ZeRO-1 / FSDP: all ranks participate in save_model (collective op) - self.latest_model = self.ckpt_store.path_for(0) - self.save_model(self.latest_model, lr=0, step=0) - if self.model_ema is not None: - self.latest_ema_model = self.ema_ckpt_store.path_for(0) - self.save_ema_model(self.latest_ema_model, lr=0, step=0) - - if ( - self.rank == 0 or dist.get_rank() == 0 - ): # Handle the case if rank 0 aborted and re-assigned - if self.num_steps == 0: - if not self.sharding.enabled: - # When num_steps is 0, the checkpoint is never saved in the loop - self.latest_model = self.ckpt_store.path_for(0) - self.save_model(self.latest_model, lr=0, step=0) - if self.model_ema is not None: - self.latest_ema_model = self.ema_ckpt_store.path_for(0) - self.save_ema_model(self.latest_ema_model, lr=0, step=0) - log.info(f"Saved model to {self.latest_model}") - self.ckpt_store.publish(self.latest_model) - if self.model_ema is not None: - self.ema_ckpt_store.publish(self.latest_ema_model) + if self.num_steps == 0: + self._save_training_checkpoint(step=0, display_step=0, lr=0.0) + if self.rank == 0 or dist.get_rank() == 0: + # Handle the case if rank 0 aborted and was reassigned. if self.timing_in_training: average_message = self.step_timer.format_average() if average_message is not None: @@ -1933,10 +1811,9 @@ def log_loss_valid(_task_key: str = "Default") -> dict: f"The profiling trace has been saved to: {self.profiling_file}" ) - def _close_lmdb_loaders(self) -> None: - """Release LMDB pipelines owned by training and validation loaders.""" + def _close_data_loaders(self) -> None: + """Release native pipelines and readers owned by both dataset splits.""" closed: set[int] = set() - datasets: dict[int, LmdbDataset] = {} for loaders in (self.training_dataloader, self.validation_dataloader): values = loaders.values() if isinstance(loaders, dict) else (loaders,) for loader in values: @@ -1946,10 +1823,32 @@ def _close_lmdb_loaders(self) -> None: close = getattr(loader, "close", None) if close is not None: close() - if isinstance(loader, LmdbBatchDataLoader): - datasets[id(loader.dataset)] = loader.dataset - for dataset in datasets.values(): - dataset.close() + + def _save_training_checkpoint( + self, *, step: int, display_step: int, lr: float + ) -> None: + """Publish live and optional EMA checkpoints after one all-rank check.""" + self.ensure_finite_gradients_for_checkpoint() + # Preserve chief ownership if rank 0 was reassigned. + is_chief = self.rank == 0 or dist.get_rank() == 0 + if not is_chief and not self.sharding.enabled: + return + self.latest_model = self.ckpt_store.path_for(display_step) + self.save_model(self.latest_model, lr=lr, step=step) + if is_chief: + log.info(f"Saved model to {self.latest_model}") + self.ckpt_store.publish(self.latest_model) + if self.model_ema is not None: + self.latest_ema_model = self.ema_ckpt_store.path_for(display_step) + self.save_ema_model(self.latest_ema_model, lr=lr, step=step) + if is_chief: + self.ema_ckpt_store.publish(self.latest_ema_model) + + def ensure_finite_gradients_for_checkpoint(self) -> None: + """Check the accumulated gradient norm on all ranks before publication.""" + self.nonfinite_grad_guard.raise_if_nonfinite( + self.wrapper.named_parameters, synchronize=self.is_distributed + ) def _collect_checkpoint_states( self, @@ -2205,7 +2104,11 @@ def get_data( return {}, {}, {} batch_data = next(iterator) # === Filter frames with atoms too close (training only) === - filter_enabled = is_train and self.min_pair_dist_by_task[task_key] > 0.0 + filter_enabled = is_train and getattr( + batch_data, + "pair_filter_enabled", + self.pair_filter_enabled_by_task[task_key], + ) valid_mask = pair_margin_frame_mask(batch_data, filter_enabled) local_has_valid = valid_mask is None or bool(torch.any(valid_mask)) if not local_has_valid: @@ -2231,6 +2134,7 @@ def get_data( "coord", "atype", "spin", + "grid", "box", "fparam", "aparam", @@ -2255,6 +2159,10 @@ def get_data( if "fid" in batch_data: log_dict["fid"] = batch_data["fid"] log_dict["sid"] = batch_data["sid"] + # models without grid support do not accept a grid keyword in forward; + # drop it when the data does not provide grid + if input_dict.get("grid") is None: + input_dict.pop("grid", None) return input_dict, label_dict, log_dict def print_header( @@ -2423,6 +2331,22 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: ) ] additional_data_requirement += spin_requirement_items + has_grid = getattr(_model, "has_grid", False) + if callable(has_grid): + has_grid = has_grid() + if has_grid: + # the grid is a model input (like fparam/aparam), not a label; it is + # frame-major (nframes, ngrid, 3) and not per-atom + additional_data_requirement.append( + DataRequirementItem( + "grid", + ndof=3, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + ) if _model.has_chg_spin_ebd(): has_default_cs = _model.has_default_chg_spin() cs_default = ( @@ -2526,6 +2450,9 @@ def get_loss( tensor_name = "polar" loss_params["tensor_name"] = tensor_name return TensorLoss(**loss_params) + elif loss_type == "grid_density": + loss_params["starter_learning_rate"] = start_lr + return GridDensityLoss(**loss_params) elif loss_type == "property": task_dim = _model.get_task_dim() var_name = _model.get_var_name() @@ -2603,6 +2530,7 @@ def model_change_out_bias( _model: Any, _sample_func: Callable[[], Any], _bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> Any: from deepmd.pt.model.model.dp_model import ( DPModelCommon, @@ -2613,4 +2541,5 @@ def model_change_out_bias( _sample_func, bias_adjust_mode=_bias_adjust_mode, recompute_input_stats=isinstance(_model, DPModelCommon), + preset_out_bias=preset_out_bias, ) diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index b2e4d4f229..7d7308b0e9 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -172,6 +172,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, ) -> tuple[Any, Any, Any]: if not self.multi_task: task_key = "Default" @@ -194,6 +195,12 @@ def forward( if has_spin: input_dict["spin"] = spin + has_grid = getattr(self.model[task_key], "has_grid", False) + if callable(has_grid): + has_grid = has_grid() + if has_grid: + input_dict["grid"] = grid + # A loss-free wrapper is a pure inference object, so parameters can be # treated as constants while coordinate gradients remain enabled. if self.inference_only: diff --git a/deepmd/pt/utils/auto_batch_size.py b/deepmd/pt/utils/auto_batch_size.py index f20b0e5f0a..d411b876ed 100644 --- a/deepmd/pt/utils/auto_batch_size.py +++ b/deepmd/pt/utils/auto_batch_size.py @@ -1,111 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Any, -) - import torch -from deepmd.pt.utils import ( - env, -) from deepmd.utils.batch_size import AutoBatchSize as AutoBatchSizeBase -from deepmd.utils.batch_size import ( - BatchSizeBudget, -) - - -class _CudaMemoryBudget: - """Estimate batch workspace without modifying allocator peak statistics. - - Calibration releases unused cache once. Reserved-memory growth together - with the remaining default-pool cache bounds batch workspace. Private CUDA - pools remain reserved for their graphs and are not reusable by evaluation. - """ - - def __init__(self, device: int, natoms: int, stats: dict[str, Any]) -> None: - self.device = device - self.natoms = natoms - self.reserved = stats["reserved_bytes"]["all"]["current"] - self.allocated = stats["allocated_bytes"]["all"]["current"] - self.allocated_total = stats["allocated_bytes"]["all"]["allocated"] - self.peak = stats["allocated_bytes"]["all"]["peak"] - self.cached = self._cached_bytes() - self.measured_atoms = 0 - self.bytes_per_atom = 0.0 - self.device_events = self._device_events(stats) - - @staticmethod - def _device_events(stats: dict[str, Any]) -> tuple[int, int]: - return stats["num_device_alloc"], stats["num_device_free"] - - def matches(self, device: int, natoms: int, stats: dict[str, Any]) -> bool: - return ( - self.device == device - and self.natoms == natoms - and self.device_events == self._device_events(stats) - ) - - def observe(self, nframes: int) -> None: - stats = torch.cuda.memory_stats_as_nested_dict(self.device) - self.device_events = self._device_events(stats) - allocated_total = stats["allocated_bytes"]["all"]["allocated"] - allocated = allocated_total - self.allocated_total - self.allocated_total = allocated_total - atoms = nframes * self.natoms - if atoms == 0 or atoms < self.measured_atoms: - return - workspace = max( - 0, stats["reserved_bytes"]["all"]["current"] - self.reserved + self.cached - ) - peak = stats["allocated_bytes"]["all"]["peak"] - if peak > self.peak: - workspace = max(workspace, peak - self.allocated) - # Unreleasable split blocks can supply workspace without growing the - # reservation or exceeding a training peak. The initial cache includes - # those blocks; cumulative batch allocations bound unused cache slack. - workspace = min(workspace, allocated) - self.measured_atoms = atoms - self.bytes_per_atom = workspace / atoms - - def _cached_bytes(self) -> int: - return sum( - segment["total_size"] - segment["active_size"] - for segment in torch.cuda.memory_snapshot() - if segment["device"] == self.device and segment["segment_pool_id"] == (0, 0) - ) - - def _available_bytes(self) -> tuple[int, int]: - """Return available device bytes and reusable default-pool cache bytes.""" - stats = torch.cuda.memory_stats_as_nested_dict(self.device) - self.device_events = self._device_events(stats) - free, total = torch.cuda.mem_get_info(self.device) - reserved = stats["reserved_bytes"]["all"]["current"] - allowed = int(total * torch.cuda.get_per_process_memory_fraction(self.device)) - return min(free, max(0, allowed - reserved)), self._cached_bytes() - - def _capacity(self, available: int) -> int: - # The unused fraction covers allocator rounding and moderate variation - # in frame density; unpredictable workspaces still use OOM backoff. - capacity = int(0.9 * available / self.bytes_per_atom) - return max(self.natoms, capacity // self.natoms * self.natoms) - - def limit(self, batch_size: int) -> int: - if not self.measured_atoms: - return min(batch_size, self.natoms) - if not self.bytes_per_atom: - return batch_size - free, cached = self._available_bytes() - target = min(batch_size, self._capacity(free + cached)) - if target > self.measured_atoms and target * self.bytes_per_atom > free: - # A larger allocation may not fit an existing cached block. Release - # unused blocks before requesting more pages, rather than relying - # on the allocator's OOM-triggered cache-flush retry. - with torch.cuda.device(self.device): - torch.cuda.empty_cache() - free, cached = self._available_bytes() - target = min(target, self._capacity(free + cached)) - return target class AutoBatchSize(AutoBatchSizeBase): @@ -121,37 +18,13 @@ class AutoBatchSize(AutoBatchSizeBase): Notes ----- - Automatic multi-frame CUDA evaluation starts with one calibration frame - and limits growth using observed workspace and available device memory. - Explicit ``DP_INFER_BATCH_SIZE`` limits and single-frame calls retain their - original behavior. Allocator peak counters and memory limits are unchanged. + The shared controller grows successful batches and remembers OOM bounds. + Successful inference neither profiles the CUDA allocator nor clears its + cache. CUDA-specific behavior is limited to recognizing and recovering + from out-of-memory errors. """ - def _get_batch_budget( - self, natoms: int, max_nframes: int | None - ) -> BatchSizeBudget | None: - if ( - self._fixed_batch_size - or max_nframes == 1 - or natoms <= 0 - or env.DEVICE.type != "cuda" - or not self.is_gpu_available() - or torch.cuda.get_allocator_backend() != "native" - ): - return None - with torch.cuda.device(env.DEVICE): - device = torch.cuda.current_device() - stats = torch.cuda.memory_stats_as_nested_dict(device) - if isinstance(self._budget, _CudaMemoryBudget) and self._budget.matches( - device, natoms, stats - ): - return self._budget - torch.cuda.empty_cache() - return _CudaMemoryBudget( - device, natoms, torch.cuda.memory_stats_as_nested_dict(device) - ) - def is_gpu_available(self) -> bool: """Check if GPU is available. diff --git a/deepmd/pt/utils/compile_compat.py b/deepmd/pt/utils/compile_compat.py index f82307b488..6f8c349cea 100644 --- a/deepmd/pt/utils/compile_compat.py +++ b/deepmd/pt/utils/compile_compat.py @@ -544,9 +544,11 @@ def build_inductor_compile_options(*, inference: bool = False) -> dict[str, Any] misbehaved on the combination of data-dependent edge counts and a second-order autograd graph -- most visibly the oversized fused Triton reduction kernels that fail ``make_ttgir`` (``PassManager::run failed``) on - some GPU/Triton combinations. Options absent from the running PyTorch's - configuration registry are dropped so the returned dictionary stays valid - across releases. + some GPU/Triton combinations. It also gives every generated kernel output + buffers separate from its inputs, without which Triton 3.7.1 can compute a + reduction wrongly. Options absent from the running PyTorch's configuration + registry are dropped so the returned dictionary stays valid across + releases. Parameters ---------- @@ -595,6 +597,18 @@ def build_inductor_compile_options(*, inference: bool = False) -> dict[str, Any] # access. A 1D grid keeps that axis on the x dimension (limit 2**31-1). # The option is shared by the training and evaluation graphs. "triton.max_tiles": 1, + # Inductor can generate a reduction kernel that loads an input inside + # its reduction loop and stores a result over that input in the same + # iteration. Where the tuned configuration launches more threads than + # the block has elements, Triton 3.7.1 replicates the block across + # warps: every copy loads, only the first copy stores, and no barrier + # separates the loads of an iteration from that store, so a lagging + # copy can load the stored value and the reduction comes out wrong + # intermittently. The Inductor cache keeps the tuned configuration, so + # every later process that loads the kernel from the cache can compute + # the reduction wrongly as well. Separate output buffers leave no + # kernel that stores over its own inputs. + "inplace_buffers": False, } # ``DP_TUNE_TRAIN`` grades the compile-time investment of the training # graphs (cumulative levels; inference graphs ignore it, the AOTI export @@ -634,6 +648,11 @@ def build_inductor_compile_options(*, inference: bool = False) -> dict[str, Any] if tune_train >= 2: compile_options["max_autotune_gemm"] = True if inference: + # PyTorch 2.11 CPU horizontal fusion can reuse a scatter's index + # array in an earlier gather, before the array is declared. Keep + # those loops separate while retaining SIMD and vertical fusion. + if _torch_release() == (2, 11): + compile_options["cpp.max_horizontal_fusion_size"] = 1 # The peak-memory reordering pass sizes buffers through # ``sizevars.size_hint(numel, fallback=0)``. The inference graph is # lowered from ``make_fx`` fake placeholders whose edge-count symbols diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index 90a3a3ca0c..297895ff24 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -1,6 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging import os +from collections.abc import ( + Iterator, +) +from functools import ( + cached_property, +) from multiprocessing.dummy import ( Pool, ) @@ -35,16 +41,38 @@ from deepmd.pt.utils.dataset import ( DeepmdDataSetForLoader, ) +from deepmd.pt.utils.env import ( + DEVICE, +) +from deepmd.pt.utils.lmdb_dataset import ( + LmdbBatchDataLoader, + LmdbDataset, +) from deepmd.pt.utils.utils import ( mix_entropy, + to_numpy_array, ) from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_config import ( + DataConfig, + data_group_configs, + resolve_data_groups, + resolve_data_sources, +) +from deepmd.utils.data_group import ( + GroupedDataset, + GroupedTrainingDataLoader, + TrainingDataLoader, +) +from deepmd.utils.data_sampling import ( + data_child_seeds, + global_batch_counts, + system_probabilities, +) from deepmd.utils.data_system import ( print_summary, - prob_sys_size_ext, - process_sys_probs, ) log = logging.getLogger(__name__) @@ -191,6 +219,23 @@ def construct_dataset(system: str) -> DeepmdDataSetForLoader: def __len__(self) -> int: return len(self.dataloaders) + @cached_property + def mean_batch_nframes(self) -> np.ndarray: + """Count frames per native batch, including tails and distributed padding.""" + return np.asarray( + [len(loader.sampler) / len(loader) for loader in self.dataloaders], + dtype=np.float64, + ) + + @cached_property + def uniform_sampling_mass(self) -> float: + """Return the summed systems-per-frame weights of the native streams.""" + return float(np.reciprocal(self.mean_batch_nframes).sum()) + + @property + def uniform_nominal_mass(self) -> float: + return self.uniform_sampling_mass + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: # log.warning(str(torch.distributed.get_rank())+" idx: "+str(idx)+" index: "+str(self.index[idx])) with torch.device("cpu"): @@ -253,21 +298,33 @@ def collate_batch(batch: list[dict[str, Any]]) -> dict[str, Any]: return result +class DpDataLoader(DataLoader): + """Own native CPU loading and persistent per-system reader cursors.""" + + def __iter__(self) -> Iterator[dict[str, Any]]: + with torch.device("cpu"): + return super().__iter__() + + def close(self) -> None: + """Release the persistent worker pool owned by this loader.""" + if self._iterator is not None: + self._iterator._shutdown_workers() + self._iterator = None + + def get_weighted_sampler( - training_data: Any, prob_style: str, sys_prob: bool = False + training_data: Any, prob_style: str | list[float], sys_prob: bool = False ) -> WeightedRandomSampler: - if sys_prob is False: - if prob_style == "prob_uniform": - prob_v = 1.0 / float(training_data.__len__()) - probs = [prob_v for ii in range(training_data.__len__())] - else: # prob_sys_size;A:B:p1;C:D:p2 or prob_sys_size = prob_sys_size;0:nsys:1.0 - if prob_style == "prob_sys_size": - style = f"prob_sys_size;0:{len(training_data)}:1.0" - else: - style = prob_style - probs = prob_sys_size_ext(style, len(training_data), training_data.index) - else: - probs = process_sys_probs(prob_style, training_data.index) + probs = system_probabilities( + training_data.index, + sys_probs=prob_style if sys_prob else None, + auto_prob="prob_sys_size" if sys_prob else prob_style, + batch_nframes=( + training_data.mean_batch_nframes + if not sys_prob and prob_style == "prob_uniform" + else None + ), + ) log.debug("Generated weighted sampler with prob array: " + str(probs)) # training_data.total_batch is the size of one epoch, you can increase it to avoid too many rebuilding of iterators len_sampler = training_data.total_batch * max(env.NUM_WORKERS, 1) @@ -294,3 +351,131 @@ def get_sampler_from_params(_data: Any, _params: dict[str, Any]) -> Any: else: _sampler = get_weighted_sampler(_data, "prob_sys_size") return _sampler + + +def create_dataset( + config: DataConfig, + type_map: list[str], + *, + seed: int | list[int] | None = None, + modifier: BaseModifier | None = None, +) -> DpLoaderSet | LmdbDataset | GroupedDataset: + """Construct native readers from the shared recursive group configuration.""" + groups = resolve_data_groups(config) + configs = data_group_configs(config) + datasets = [] + for group, params in zip(groups, configs, strict=True): + sources = resolve_data_sources(group) + has_lmdb = any(source.format == "lmdb" for source in sources) + children = [] + child_configs = [] + for source in sources: + if source.format == "deepmd": + child = DpLoaderSet( + list(source.systems), + source.batch_size, + type_map, + seed=seed, + modifier=modifier, + ) + else: + child = LmdbDataset( + source.systems[0], + type_map, + source.batch_size, + auto_prob_style=group.auto_prob, + ) + children.append(child) + child_configs.append( + {**params, "sys_probs": group.sys_probs if not has_lmdb else None} + ) + datasets.append( + children[0] + if len(children) == 1 + else GroupedDataset( + children, + child_configs, + [None] * len(children), + uniform=group.auto_prob == "prob_uniform", + ) + ) + return ( + datasets[0] + if len(datasets) == 1 + else GroupedDataset(datasets, configs, [group.prob for group in groups]) + ) + + +def create_dataloader( + data: DpLoaderSet | LmdbDataset | GroupedDataset, + config: DataConfig, + *, + seed: int | None = None, + rank: int = 0, + world_size: int = 1, +) -> TrainingDataLoader | GroupedTrainingDataLoader: + """Bind the native batch pipelines and their local sampling metadata.""" + if isinstance(data, GroupedDataset): + seeds = data_child_seeds( + 0 if seed is None and world_size > 1 else seed, len(data.children) + 1 + ) + children = [ + create_dataloader( + child, params, seed=child_seed, rank=rank, world_size=world_size + ) + for child, params, child_seed in zip( + data.children, data.configs, seeds[1:], strict=True + ) + ] + return GroupedTrainingDataLoader(data, children, seeds[0]) + if isinstance(data, LmdbDataset): + from deepmd.dpmodel.utils.lmdb_data import ( + DistributedLmdbBatchSampler, + LmdbBatchSampler, + ) + + options = {"shuffle": True, "seed": seed, "block_targets": data._block_targets} + sampler = ( + DistributedLmdbBatchSampler( + data._reader, rank=rank, world_size=world_size, **options + ) + if world_size > 1 + else LmdbBatchSampler(data._reader, **options) + ) + loader = LmdbBatchDataLoader(data, sampler, pin_memory=DEVICE.type != "cpu") + return TrainingDataLoader( + data, + loader, + np.ones(1, dtype=np.float64), + lambda: [len(loader)], + sampling_sizes=lambda: [sampler.total_batches], + uniform_sampling=loader._batch_iterator, + ) + params = config[0] if isinstance(config, list) else config + sampler = get_sampler_from_params(data, params) + num_workers = env.NUM_WORKERS if dist.is_available() else 0 + loader = DpDataLoader( + data, + sampler=sampler, + batch_size=None, + num_workers=num_workers, + persistent_workers=( + num_workers > 0 + and params.get("auto_prob") == "prob_uniform" + and params.get("sys_probs") is None + ), + drop_last=False, + collate_fn=lambda batch: batch, + pin_memory=DEVICE.type != "cpu", + ) + + return TrainingDataLoader( + data, + loader, + to_numpy_array(sampler.weights), + lambda: data.index, + sampling_sizes=lambda: global_batch_counts( + [len(system) for system in data.systems], data.batch_sizes + ), + uniform_sampling=data, + ) diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py index d3a3583938..616d9cae9a 100644 --- a/deepmd/pt/utils/dataset.py +++ b/deepmd/pt/utils/dataset.py @@ -74,6 +74,7 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N default=data_item["default"], dtype=data_item["dtype"], output_natoms_for_type_sel=data_item["output_natoms_for_type_sel"], + special_shape=getattr(data_item, "special_shape", None), length_scale=data_item["length_scale"], ) diff --git a/deepmd/pt/utils/finetune.py b/deepmd/pt/utils/finetune.py index bf1c328408..8a46ae0d69 100644 --- a/deepmd/pt/utils/finetune.py +++ b/deepmd/pt/utils/finetune.py @@ -21,6 +21,33 @@ __all__ = ["get_finetune_rule_single", "get_finetune_rules"] +def _learned_part_path(model: torch.nn.Module) -> str | None: + """Return the state-dict path of the learned descriptor-fitting part of a model. + + A plain model holds that part as its atomic model and a composition with + one learned sub-model, such as a zone-bridging one, as that sub-model, so + fine-tuning between the two transfers the part across the paths. A model + without a single learned part returns ``None``. + """ + atomic_model = model.atomic_model + if getattr(atomic_model, "descriptor", None) is not None: + return "atomic_model." + learned = [ + index + for index, child in enumerate(getattr(atomic_model, "models", ())) + if getattr(child, "descriptor", None) is not None + ] + return f"atomic_model.models.{learned[0]}." if len(learned) == 1 else None + + +def _move_learned_key(key: str, path: str, source_path: str) -> str: + """Move a descriptor, fitting-net or long-range key between learned atomic-model paths.""" + for part in ("descriptor.", "fitting_net.", "long_range."): + if f".{path}{part}" in key: + return key.replace(f".{path}{part}", f".{source_path}{part}", 1) + return key + + def get_finetune_rules( finetune_model: str, model_config: dict[str, Any], diff --git a/deepmd/pt/utils/lmdb_dataset.py b/deepmd/pt/utils/lmdb_dataset.py index 8cce38a853..70dd25b864 100644 --- a/deepmd/pt/utils/lmdb_dataset.py +++ b/deepmd/pt/utils/lmdb_dataset.py @@ -27,10 +27,8 @@ collate_lmdb_frames, collect_lmdb_sampling_groups, compute_block_targets, - count_group_blocks, is_lmdb, resolve_per_atom_keys, - system_block_lookup, ) from deepmd.env import ( get_lmdb_num_workers, @@ -347,14 +345,12 @@ def print_summary(self, name: str, prob: Any) -> None: reader = self._reader # Per-block summary: original vs target frames block_lines = [] - total_original = 0 + total_original = len(reader) total_target = 0 - # Pre-compute block_total_actual for proportional scaling - block_total_actual: list[int] = [] + assigned_systems: set[int] = set() for sys_ids, target in self._block_targets: actual = sum(reader.system_nframes[s] for s in sys_ids) - block_total_actual.append(actual) - total_original += actual + assigned_systems.update(sys_ids) total_target += target # Compact range notation: sys[0-146] instead of sys[0,1,2,...,146] if len(sys_ids) > 3: @@ -367,42 +363,18 @@ def print_summary(self, name: str, prob: Any) -> None: f"{actual}->{target} (x{ratio:.2f})" ) - # A whole nloc group's block membership resolves as one indexing - # operation, which a dataset of 10^8 frames needs it to be. The - # lookup is the one the sampler allocates its targets with, so - # both agree on which systems a block claims. - n_blocks = len(self._block_targets) - lookup = system_block_lookup(self._block_targets) - - # Compute expanded nloc counts analytically (no actual expansion) - expanded_nloc_info = {} - for nloc, indices in sorted(reader.nloc_groups.items()): - if reader.frame_system_ids is None: - expanded_nloc_info[nloc] = len(indices) - continue - counts = count_group_blocks( - indices, reader.frame_system_ids, lookup, n_blocks - ) - expanded = len(indices) - int(counts.sum()) - for blk_idx, (_, blk_target) in enumerate(self._block_targets): - n_actual = int(counts[blk_idx]) - if n_actual == 0: - continue - bta = block_total_actual[blk_idx] - if bta > 0: - t = max(round(blk_target * n_actual / bta), n_actual) - else: - t = n_actual - expanded += t - expanded_nloc_info[nloc] = expanded - - total_expanded = sum(expanded_nloc_info.values()) + # Systems outside explicit blocks retain their original frames. + total_target += sum( + count + for sid, count in enumerate(reader.system_nframes) + if sid not in assigned_systems + ) n_groups = len(reader.nloc_groups) - ratio_all = total_expanded / total_original if total_original > 0 else 0 + ratio_all = total_target / total_original if total_original > 0 else 0 log.info( f"LMDB {name} auto_prob: " - f"{total_original}->{total_expanded} frames (x{ratio_all:.2f}), " + f"{total_original}->{total_target} frames (x{ratio_all:.2f}), " f"{n_groups} nloc groups, {len(self._block_targets)} blocks:" ) for bl in block_lines: diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index 3c68251a7d..cc3b925afb 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -172,9 +172,7 @@ def build( When provided, excluded type pairs are erased from the returned neighbor list (entries set to ``-1``) by :func:`~deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist`. - ``NvNeighborList`` is CUDA-only; the ``pair_excl`` parameter is - accepted for API parity with the other strategies but cannot be - validated on a CPU-only machine. + Search and exclusion run on the input device, either CPU or CUDA. ``return_mode='edges'`` does not support ``pair_excl``; a :class:`NotImplementedError` is raised in that combination. """ diff --git a/deepmd/pt/utils/serialization.py b/deepmd/pt/utils/serialization.py index 6cdd65712f..a8a5681ef9 100644 --- a/deepmd/pt/utils/serialization.py +++ b/deepmd/pt/utils/serialization.py @@ -80,12 +80,6 @@ def deserialize_to_file(model_file: str, data: dict) -> None: ) model = SpinEnergyModel.deserialize(model_data) - elif model_data.get("type") == "sezm_spin": - from deepmd.pt.model.model.sezm_spin_model import ( - SeZMSpinModel, - ) - - model = SeZMSpinModel.deserialize(model_data) elif model_data.get("type") == "sezm_native_spin": from deepmd.pt.model.model.sezm_native_spin_model import ( SeZMNativeSpinModel, diff --git a/deepmd/pt/utils/spin_finetune.py b/deepmd/pt/utils/spin_finetune.py new file mode 100644 index 0000000000..064eed85df --- /dev/null +++ b/deepmd/pt/utils/spin_finetune.py @@ -0,0 +1,342 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Spin-free checkpoint transfer for the PyTorch and pt_expt native-spin models.""" + +from collections.abc import ( + Callable, +) +from functools import ( + singledispatch, +) + +import torch + +from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, +) +from deepmd.dpmodel.fitting.dpa4_ener import ( + GLUFittingNet, +) +from deepmd.dpmodel.utils.network import ( + FittingNet, + NativeLayer, +) +from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, +) +from deepmd.pt.utils.finetune import ( + _learned_part_path, +) + + +def get_native_spin_finetune_state( + model: torch.nn.Module, + pretrained_model: torch.nn.Module, + *, + random_fitting: bool = False, +) -> dict[str, torch.Tensor]: + """Build native-spin state from a spin-free learned atomic model. + + Type maps must already agree. The learned parts are matched independently + of analytical bridging. Output biases and scales belong to the target's + subsequent statistics pass; spin-only parameters keep their constructor + values, while target data calibrates spin-only input statistics afterward. + Descriptor-specific transfer rules preserve the geometric function while + leaving a nonzero gradient into the initially closed spin routes. + + Parameters + ---------- + model + Fresh native-spin target, with its input configuration and spin mask. + pretrained_model + Loaded spin-free source after type-map alignment. + random_fitting + Whether to retain the target fitting network instead of transferring it. + + Returns + ------- + dict[str, torch.Tensor] + Complete state dict for a strict load into the target model. + + Raises + ------ + ValueError + If the learned parts, descriptor or fitting families, type maps, + network layouts or tensor shapes are incompatible with spin-free transfer, + or widening the fitting input changes an implicit residual connection. + """ + path = _learned_part_path(model) + source_path = _learned_part_path(pretrained_model) + if path is None or source_path is None: + raise ValueError( + "Spin-free transfer requires one learned atomic model in each model." + ) + if model.get_type_map() != pretrained_model.get_type_map(): + raise ValueError("Spin-free transfer requires aligned type maps.") + learned = model.get_submodule(path.removesuffix(".")) + source = pretrained_model.get_submodule(source_path.removesuffix(".")) + descriptor = learned.descriptor + if type(descriptor) is not type(source.descriptor): + raise ValueError( + "Spin-free transfer requires the same descriptor family: " + f"source={type(source.descriptor).__name__}, " + f"target={type(descriptor).__name__}." + ) + if descriptor.use_spin is None or source.descriptor.use_spin is not None: + raise ValueError( + "Spin-free transfer requires a spin-free source descriptor and a " + "native-spin target descriptor." + ) + + expanded = _expanded_spin_state( + descriptor, source, learned, random_fitting=random_fitting + ) + if not random_fitting: + _validate_fitting_networks(source.fitting_net, learned.fitting_net) + state = model.state_dict() + for key, value in source.state_dict().items(): + target_key = path + key + if target_key not in state or key.rsplit(".", 1)[-1] in { + "out_bias", + "out_std", + "bias_atom_e", + }: + continue + if random_fitting and not key.startswith("descriptor."): + continue + if key in expanded: + state[target_key] = expanded[key] + else: + _check_shape(key, value.shape, state[target_key].shape) + state[target_key] = value.detach().clone() + return state + + +def compute_native_spin_input_stats( + model: torch.nn.Module, sampled_func: Callable[[], list[dict]] +) -> None: + """Calibrate newly introduced spin state after spin-free weight transfer. + + Geometric input statistics remain those of the pretrained descriptor. + Only DPA4C requires a data-dependent reference for its spin invariants. + + Parameters + ---------- + model + Native-spin target after a successful spin-free state transfer. + sampled_func + Sampled training systems containing atom types and magnetic moments. + """ + path = _learned_part_path(model) + descriptor = model.get_submodule(path.removesuffix(".")).descriptor + if isinstance(descriptor, DescrptDPA4C): + descriptor.spin.set_spin_reference( + descriptor._measure_spin_reference(sampled_func()) + ) + + +@singledispatch +def _expanded_spin_state( + descriptor: object, + source: torch.nn.Module, + target: torch.nn.Module, + *, + random_fitting: bool, +) -> dict[str, torch.Tensor]: + """Expand state under the descriptor's native-spin layout contract.""" + raise ValueError( + f"Spin-free transfer is not registered for {type(descriptor).__name__}." + ) + + +@_expanded_spin_state.register(DescrptSeZM) +@_expanded_spin_state.register(DescrptDPA4) +def _expanded_dpa4_state( + descriptor: object, + source: torch.nn.Module, + target: torch.nn.Module, + *, + random_fitting: bool, +) -> dict[str, torch.Tensor]: + # DPA4 injects zero-initialized additive spin routes without widening the + # geometric descriptor or its fitting input. + return {} + + +@_expanded_spin_state.register(DescrptDPA4C) +def _expanded_dpa4c_state( + descriptor: DescrptDPA4C, + source: torch.nn.Module, + target: torch.nn.Module, + *, + random_fitting: bool, +) -> dict[str, torch.Tensor]: + # The descriptor is [readout, spin, masses, type embedding]. Any frame or + # atomic fitting parameters follow this descriptor and shift with its tail. + start = descriptor.readout.get_dim_out() + width = descriptor.spin.get_dim_out() + expanded = ( + {} + if random_fitting + else _expanded_fitting_input_state( + source.fitting_net, target.fitting_net, start=start, width=width + ) + ) + source_state = source.state_dict() + target_state = target.state_dict() + for key in ("descriptor.mean", "descriptor.stddev"): + expanded[key] = _insert_spin_block( + key, + source_state[key], + target_state[key], + axis=0, + start=start, + width=width, + ) + # Ordered spin scale/shift columns follow every geometric FiLM head. + key = "descriptor.pair_film.network.layers.1.w" + expanded[key] = _insert_spin_block( + key, + source_state[key], + target_state[key], + axis=1, + start=source_state[key].shape[1], + width=2 * descriptor.spin_channels, + ) + return expanded + + +@singledispatch +def _fitting_input_projection(network: object) -> NativeLayer: + """Return the affine layer that reads the fitting input.""" + raise ValueError( + "Spin-free input expansion is not registered for fitting network " + f"{type(network).__name__}." + ) + + +@_fitting_input_projection.register(FittingNet) +def _mlp_input_projection(network: FittingNet) -> NativeLayer: + return network.layers[0] + + +@_fitting_input_projection.register(GLUFittingNet) +def _glu_input_projection(network: GLUFittingNet) -> NativeLayer: + return network.hidden_layers[0] if network.hidden_layers else network.output_layer + + +def _expanded_fitting_input_state( + source: torch.nn.Module, + target: torch.nn.Module, + *, + start: int, + width: int, +) -> dict[str, torch.Tensor]: + """Expand fitting input projections without changing their residual semantics.""" + if source.nets.ndim != target.nets.ndim: + raise ValueError("Spin-free transfer requires the same fitting type layout.") + expanded: dict[int, torch.Tensor] = {} + for index in range(target.nets.ntypes**target.nets.ndim): + source_net, target_net = source.nets[index], target.nets[index] + if type(source_net) is not type(target_net): + raise ValueError( + "Spin-free transfer requires matching fitting network types: " + f"source={type(source_net).__name__}, " + f"target={type(target_net).__name__}." + ) + source_layer = _fitting_input_projection(source_net) + target_layer = _fitting_input_projection(target_net) + # NativeLayer derives its residual and timestep from the input/output + # widths. An inserted input block changes an identity or duplicated + # residual, which cannot in general be absorbed into the nonlinear + # branch's weights. Reject that structural change before loading state. + for layer in (source_layer, target_layer): + if layer.resnet and layer.dim_out() in ( + layer.dim_in(), + 2 * layer.dim_in(), + ): + raise ValueError( + "Spin-free transfer cannot preserve the first fitting layer's " + "residual connection when its input expands from " + f"{source_layer.dim_in()} to {target_layer.dim_in()} " + f"with output width {target_layer.dim_out()}. " + "Use a compatible pretrained fitting architecture or select " + "finetune_head='RANDOM' to initialize a fresh fitting network." + ) + # A closed descriptor spin gate preserves the geometric function while + # fresh fitting rows give that gate a nonzero gradient. + expanded[id(target_layer)] = _insert_spin_block( + f"fitting_net.nets[{index}].input_projection.w", + source_layer.w, + target_layer.w, + axis=0, + start=start, + width=width, + ) + # Network collections expose multiple state-dict paths for the same layer. + # Every alias refers to the same expanded weight, computed once per network. + return { + f"fitting_net.{name}.w": expanded[id(module)] + for name, module in target.named_modules(remove_duplicate=False) + if id(module) in expanded + } + + +def _validate_fitting_networks( + source: torch.nn.Module, target: torch.nn.Module +) -> None: + """Require complete fitting-network transfer, including frozen weights.""" + if type(source) is not type(target): + raise ValueError( + "Spin-free transfer requires matching fitting network types: " + f"source={type(source).__name__}, target={type(target).__name__}." + ) + # PT and pt_expt register their networks under these respective paths. + # Root buffers hold output biases, statistics and runtime defaults, not + # the learned network layout. Frozen network weights remain in state_dict. + source_keys, target_keys = ( + { + key + for key in fitting.state_dict() + if key.startswith(("filter_layers.", "nets.")) + } + for fitting in (source, target) + ) + if source_keys != target_keys: + raise ValueError( + "Spin-free transfer requires matching fitting network state keys: " + f"source-only={sorted(source_keys - target_keys)}, " + f"target-only={sorted(target_keys - source_keys)}. " + "Use a compatible fitting architecture or select finetune_head='RANDOM'." + ) + + +def _check_shape(key: str, expected: tuple[int, ...], actual: tuple[int, ...]) -> None: + if actual != expected: + raise ValueError( + f"Incompatible spin-free transfer for {key}: " + f"expected target shape {tuple(expected)}, got {tuple(actual)}." + ) + + +def _insert_spin_block( + key: str, + source: torch.Tensor, + target: torch.Tensor, + *, + axis: int, + start: int, + width: int, +) -> torch.Tensor: + """Copy the source around an inserted block, retaining its target values.""" + expected = list(source.shape) + expected[axis] += width + _check_shape(key, tuple(expected), target.shape) + source = source.detach() + result = target.detach().clone() + result.narrow(axis, 0, start).copy_(source.narrow(axis, 0, start)) + tail = source.shape[axis] - start + result.narrow(axis, start + width, tail).copy_(source.narrow(axis, start, tail)) + return result diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 49ffd995cf..3793c30bb2 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -156,6 +156,8 @@ def make_stat_input( dataloaders: list[Any], nbatches: int, min_pair_dist: float = 0.0, + *, + filter_enabled_by_system: Sequence[bool] | None = None, ) -> list[dict[str, Any]]: """Pack data for statistics. @@ -172,6 +174,9 @@ def make_stat_input( positive, frames holding a pair inside its window are dropped before the statistics are accumulated, so they see the frames the optimizer sees. + filter_enabled_by_system + Optional source-specific filter policies for grouped datasets. The + registered requirements retain each source's filter radius. Returns ------- @@ -182,12 +187,17 @@ def make_stat_input( lst = [] log.info("Packing data for statistics from %d systems", len(datasets)) for system_index in range(len(datasets)): + filter_enabled = ( + min_pair_dist > 0.0 + if filter_enabled_by_system is None + else filter_enabled_by_system[system_index] + ) sys_stat = {} with torch.device("cpu"): dataloader = dataloaders[system_index] dataloader_size = len(dataloader) target_batches = min(nbatches, dataloader_size) - scan_limit = dataloader_size if min_pair_dist > 0.0 else target_batches + scan_limit = dataloader_size if filter_enabled else target_batches iterator = iter(dataloader) accepted_batches = 0 scanned_batches = 0 @@ -201,7 +211,7 @@ def make_stat_input( except StopIteration: break scanned_batches += 1 - frame_mask = pair_margin_frame_mask(stat_data, min_pair_dist > 0.0) + frame_mask = pair_margin_frame_mask(stat_data, filter_enabled) if frame_mask is not None and not torch.any(frame_mask): continue if frame_mask is not None: @@ -279,6 +289,8 @@ def scan_redu_stats( keys: Sequence[str], intensive: bool = False, min_pair_dist: float = 0.0, + *, + filter_enabled_by_system: Sequence[bool] | None = None, ) -> ReduScanResult: """Accumulate exact reduced-label statistics over every training frame. @@ -301,6 +313,9 @@ def scan_redu_stats( min_pair_dist Radius the training-frame filter was registered with. When positive, frames holding a pair inside its window are excluded. + filter_enabled_by_system + Optional source-specific filter policies for grouped datasets. The + registered requirements retain each source's filter radius. Returns ------- @@ -319,9 +334,14 @@ def scan_redu_stats( "Scanning all frames of %d systems for output statistics", len(dataloaders) ) with torch.device("cpu"): - for dataloader in dataloaders: + for system_index, dataloader in enumerate(dataloaders): + filter_enabled = ( + min_pair_dist > 0.0 + if filter_enabled_by_system is None + else filter_enabled_by_system[system_index] + ) for batch in _full_pass_loader(dataloader): - frame_mask = pair_margin_frame_mask(batch, min_pair_dist > 0.0) + frame_mask = pair_margin_frame_mask(batch, filter_enabled) if frame_mask is not None: if not torch.any(frame_mask): continue @@ -400,11 +420,22 @@ def _compute_model_predict( charge_spin = system.get("charge_spin", None) spin = system.get("model_spin", system.get("spin", None)) + size_proxy = system["atype"].shape[-1] + if "grid" in system: + # the directional neighbor list is dense in ngrid x nall, so the + # batching size proxy must account for the grid extent. Note: no + # density model reaches this out-stat path today (out-stat is a + # documented no-op for grid models, see issue #6029), so this is + # for a future out-stat design + grid = system["grid"] + ngrid = grid.shape[-2] if grid.ndim >= 3 else grid.shape[-1] // 3 + size_proxy = max(size_proxy, ngrid) + def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: return auto_batch_size.execute_all( model_forward, nframes, - system["atype"].shape[-1], + size_proxy, *args, **kwargs, ) @@ -414,6 +445,8 @@ def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: "aparam": aparam, "charge_spin": charge_spin, } + if "grid" in system: + model_kwargs["grid"] = system["grid"] if spin is not None: model_kwargs["spin"] = spin sample_predict = model_forward_auto_batch_size( diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py index 1c19f7b1b0..b91c24d448 100644 --- a/deepmd/pt_expt/common.py +++ b/deepmd/pt_expt/common.py @@ -273,6 +273,28 @@ def _try_convert_list(name: str, value: list) -> torch.nn.Module | None: return None +def _module_buffer_device(obj: torch.nn.Module, name: str) -> torch.device: + """Device for a numpy array assigned as a buffer on ``obj``. + + Prefer the existing buffer of the same name when it already holds a + tensor, then any parameter or non-None buffer on the module, then + ``env.DEVICE``. Assignments after ``Module.to`` therefore keep residency + with the module rather than jumping back to the process default device. + """ + from deepmd.pt_expt.utils import env # deferred - avoids circular import + + if "_buffers" in obj.__dict__: + existing = obj._buffers.get(name) + if existing is not None: + return existing.device + for tensor in obj.parameters(): + return tensor.device + for tensor in obj.buffers(): + if tensor is not None: + return tensor.device + return env.DEVICE + + def dpmodel_setattr(obj: torch.nn.Module, name: str, value: Any) -> tuple[bool, Any]: """Common __setattr__ logic for pt_expt wrappers around dpmodel classes. @@ -319,6 +341,10 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer - In dpmodel, all persistent arrays use `self.xxx = np.array(...)`. Scalars use `.item()`, lists use `.tolist()`. So `isinstance(value, np.ndarray)` reliably identifies buffer-worthy attributes. + - A numpy array becomes a buffer on the module's current device (see + :func:`_module_buffer_device`), not unconditionally on ``env.DEVICE``, so + a later assignment after ``Module.to`` does not pull tensors back to the + process default device. - torch.Tensor values assigned to existing buffers fall through to torch.nn.Module.__setattr__, which correctly updates them. - dpmodel objects are identified by registry lookup (exact type match), so @@ -328,10 +354,10 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer **Circular import resolution:** - The function uses a deferred import `from deepmd.pt_expt.utils import env` - inside the function body. This breaks the circular dependency chain: - common.py → utils/__init__.py → exclude_mask.py → common.py. The import is - cached by Python after the first call, so there's no performance penalty. + Device resolution and related helpers defer + ``from deepmd.pt_expt.utils import env`` so this module does not import + ``utils`` at load time (common.py → utils/__init__.py → exclude_mask.py → + common.py). The import is cached after the first call. **Usage pattern:** @@ -354,11 +380,9 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer >>> # Case 3: dpmodel object → pt_expt module >>> obj.emask = AtomExcludeMaskDP(...) # becomes AtomExcludeMask module """ - from deepmd.pt_expt.utils import env # deferred - avoids circular import - # numpy array → torch buffer if isinstance(value, np.ndarray) and "_buffers" in obj.__dict__: - tensor = torch.as_tensor(value, device=env.DEVICE) + tensor = torch.as_tensor(value, device=_module_buffer_device(obj, name)) if name in obj._buffers: obj._buffers[name] = tensor return True, tensor diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py index a6482fbaa6..bc269c6a89 100644 --- a/deepmd/pt_expt/descriptor/dpa4c.py +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -158,6 +158,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.use_amp_infer = use_amp_infer() self._apply_autocast_policy() + def _load_from_state_dict( + self, + state_dict: dict[str, Any], + prefix: str, + *args: Any, + **kwargs: Any, + ) -> None: + # Back-compat: checkpoints written by DeePMD-kit 3.2 hold the gather + # tables of the angular moments, which are formed by broadcasting; + # the stale entries are dropped so that the strict load succeeds. + for name in ("angular_channel_index", "angular_harmonic_index"): + state_dict.pop(prefix + name, None) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + def call_graph( self, graph: Any, @@ -439,8 +453,8 @@ def deserialize(cls, data: dict) -> "DescrptDPA4C": if compression is not None: obj._set_compression( { - name: torch.as_tensor(compression["@variables"][name]) - for name in obj._COMPRESSION_BUFFER_NAMES + name: torch.as_tensor(value) + for name, value in obj._compression_variables(compression).items() } ) return obj diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 9aeae29ca3..00a5cc219c 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -1,13 +1,21 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Compress a pt_expt model (.pte) by tabulating embedding nets.""" +"""Build and export compressed pt_expt inference models.""" import logging +from typing import ( + TYPE_CHECKING, +) from deepmd.pt_expt.utils.serialization import ( deserialize_to_file, serialize_from_file, ) +if TYPE_CHECKING: + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + log = logging.getLogger(__name__) @@ -19,14 +27,14 @@ def enable_compression( check_frequency: int = -1, training_script: str | None = None, ) -> None: - """Compress a .pte model by tabulating embedding nets. + """Compress a .pte or .pt2 model by tabulating embedding nets. Parameters ---------- input_file : str - Path to the input .pte model file. + Path to the input .pte or .pt2 model file. output : str - Path to the output compressed .pte model file. + Path to the output compressed .pte or .pt2 model file. stride : float The uniform stride of the first table. extrapolate : int @@ -41,16 +49,60 @@ def enable_compression( BaseModel, ) - # 1. Load the .pte model model_dict = serialize_from_file(input_file) model = BaseModel.deserialize(model_dict["model"]) + compress_model( + model, + output, + stride=stride, + extrapolate=extrapolate, + check_frequency=check_frequency, + training_script=training_script, + min_nbor_dist=model_dict.get("min_nbor_dist"), + model_def_script=model_dict.get("model_def_script"), + ) - # 2. Get or compute min_nbor_dist. Measuring it is a dense all-pairs pass - # over the training data, so it is only run for models that tabulate - # from the shortest observed distance. - min_nbor_dist = model.get_min_nbor_dist() - if min_nbor_dist is None: - min_nbor_dist = model_dict.get("min_nbor_dist") + +def compress_model( + model: "BaseModel", + output: str, + *, + stride: float = 0.01, + extrapolate: int = 5, + check_frequency: int = -1, + training_script: str | None = None, + min_nbor_dist: float | None = None, + model_def_script: dict | None = None, + with_comm: bool = True, +) -> None: + """Compress a loaded model and export its inference artifact. + + Parameters + ---------- + model + Model whose parameters and output statistics are loaded. + output + Destination ``.pte`` or ``.pt2`` file. + stride + Uniform spacing of the first compression table. + extrapolate + Scale of the extrapolation region for models that require it. + check_frequency + Frequency of table overflow checks. + training_script + Training configuration used when neighbor statistics are required + but absent from the model and ``min_nbor_dist``. + min_nbor_dist + Fallback minimum neighbor distance from artifact metadata. + model_def_script + Training model configuration to preserve in the exported artifact. + with_comm + Whether to include a communication artifact for cross-rank models. + """ + # === Step 1. Resolve the domain required by the compression tables. === + model_min_nbor_dist = model.get_min_nbor_dist() + if model_min_nbor_dist is not None: + min_nbor_dist = model_min_nbor_dist if min_nbor_dist is None and not model.compression_needs_min_nbor_dist(): log.info( "The model tabulates over an analytically bounded domain; " @@ -64,9 +116,10 @@ def enable_compression( ) if training_script is None: raise ValueError( - "The model does not have a minimum neighbor distance, " - "so the training script and data must be provided " - "(via -t,--training-script)." + "The model does not have a minimum neighbor distance. " + "Provide training_script to compression, or freeze without " + "--compress and then run 'dp --pt-expt compress -i model.pte " + "-o compressed.pte -t INPUT.json' with the training script and data." ) from deepmd.common import ( j_loader, @@ -95,8 +148,12 @@ def enable_compression( model.min_nbor_dist = min_nbor_dist - # 3. Enable compression (also ensures fake ops are registered now that - # the C++ custom op library is loaded via enable_compression imports) + # === Step 2. Fold deployment constants and build the compressed snapshot. === + # The neutral vacuum reference must be resolved before a charge-conditioned + # descriptor specializes its tables to the deployment charge state. + model.eval() + model.fold_vacuum_reference() + uncompressed_model_dict = model.serialize() from deepmd.pt_expt.utils.tabulate_ops import ( ensure_fake_registered, ) @@ -111,16 +168,14 @@ def enable_compression( check_frequency, ) - # 4. Serialize the compressed model dict (includes tabulated data) compressed_model_dict = model.serialize() - # 5. Re-export the compressed model. + # === Step 3. Export the graph supported by the compressed model. === # - # A geometrically compressed graph-lower descriptor (DPA1 / se_atten strip, - # attn_layer == 0) evaluates its tabulated embedding through fused CUDA - # operators that make_fx can trace, so its compressed graph exports directly - # to a graph-lower ``.pt2``: ``deserialize_to_file`` bakes the end-to-end - # fused table operator and the mandatory per-atom virial (see its docstring). + # Compressed graph descriptors use fused operators with traceable fake + # implementations. Their graphs export directly to ``.pt2``; the exporter + # selects a compact canonical schema when the target supports it and + # preserves the per-atom virial alongside the total energy and forces. # Every other compressed descriptor keeps tabulated operators make_fx cannot # trace: those export the UNCOMPRESSED graph and carry the compressed dict in # ``model.json`` so ``deserialize()`` restores the compression state for the @@ -129,23 +184,24 @@ def enable_compression( model_uses_graph_lower, ) - model_def_script = model_dict.get("model_def_script") if output.endswith(".pt2") and model_uses_graph_lower(model): log.info("Re-exporting compressed graph...") deserialize_to_file( output, {"model": compressed_model_dict, "model_def_script": model_def_script}, lower_kind="auto", + with_comm=with_comm, ) else: log.info("Re-exporting compressed model...") deserialize_to_file( output, - {"model": model_dict["model"], "model_def_script": model_def_script}, + {"model": uncompressed_model_dict, "model_def_script": model_def_script}, model_json_override={ "model": compressed_model_dict, "model_def_script": model_def_script, "min_nbor_dist": float(min_nbor_dist), }, + with_comm=with_comm, ) log.info("Compressed model saved to %s", output) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index e4ff2f813d..72de9277da 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -25,9 +25,6 @@ make_task_maps, print_data_summaries, ) -from deepmd.dpmodel.utils.lmdb_data import ( - is_lmdb, -) from deepmd.loggers import ( is_node_main_process, ) @@ -37,13 +34,20 @@ from deepmd.pt_expt.utils.lmdb_dataset import ( LmdbDataSystem, ) +from deepmd.utils import random as dp_random from deepmd.utils.data import ( min_pair_dist_requirement, ) +from deepmd.utils.data_config import ( + DataConfig, +) +from deepmd.utils.data_path import ( + process_systems, +) from deepmd.utils.data_system import ( DeepmdDataSystem, get_data, - process_systems, + get_neighbor_stat_data, ) from deepmd.utils.stat_file import ( StatFileSpec, @@ -93,87 +97,32 @@ def _update_changed_model_tensors( target_state_dict[key] = source_value -def _detect_lmdb_path(systems_raw: Any) -> str | None: - """Return the LMDB path when ``systems_raw`` is a scalar LMDB string. - - Returns ``None`` for non-LMDB inputs. Raises ``ValueError`` if - ``systems_raw`` is a list containing any LMDB path, so both - ``_get_neighbor_stat_data`` and ``_build_data_system`` fail with the - same clear message instead of the opaque error from - :func:`process_systems` / :class:`DeepmdData`. - """ - if isinstance(systems_raw, str) and is_lmdb(systems_raw): - return systems_raw - if isinstance(systems_raw, list) and any( - isinstance(s, str) and is_lmdb(s) for s in systems_raw - ): - raise ValueError( - "LMDB datasets must be passed as a scalar 'systems' string " - "(e.g. 'systems': '/path/to/data.lmdb'); list-form systems " - "with LMDB paths are not supported." - ) - return None - - -def _get_neighbor_stat_data( - dataset_params: dict[str, Any], - type_map: list[str] | None, -) -> Any: - """Return a data proxy suitable for ``BaseModel.update_sel`` (neighbor stat). - - Routes a scalar LMDB ``systems`` path through dpmodel's - ``make_neighbor_stat_data``; falls back to the legacy ``get_data`` for - npy/HDF5 directories. - """ - lmdb_path = _detect_lmdb_path(dataset_params.get("systems")) - if lmdb_path is not None: - from deepmd.dpmodel.utils.lmdb_data import ( - make_neighbor_stat_data, - ) - - return make_neighbor_stat_data(lmdb_path, type_map) - return get_data(dataset_params, 0, type_map, None) - - def _build_data_system( - dataset_params: dict[str, Any], + dataset_params: DataConfig, type_map: list[str], seed: int | None = None, rank: int = 0, world_size: int = 1, -) -> DeepmdDataSystem | LmdbDataSystem: - """Build a data system from dataset config, routing LMDB paths to LmdbDataSystem. - - A scalar ``systems`` value pointing at an LMDB directory triggers the - LMDB adapter; otherwise we fall through to the legacy - :class:`DeepmdDataSystem` path with system expansion. ``rank`` and - ``world_size`` shard LMDB training batches without changing legacy data - systems. - """ - systems_raw = dataset_params["systems"] - lmdb_path = _detect_lmdb_path(systems_raw) - if lmdb_path is not None: +) -> Any: + """Bind native LMDB batching to the common grouped-data factory.""" + if seed is None and world_size > 1: + seed = 0 + + def build_lmdb( + path: str, batch_size: int | str, auto_prob: str | None + ) -> LmdbDataSystem: return LmdbDataSystem( - lmdb_path=lmdb_path, + lmdb_path=path, type_map=type_map, - batch_size=dataset_params["batch_size"], - auto_prob_style=dataset_params.get("auto_prob"), + batch_size=batch_size, + auto_prob_style=auto_prob, seed=seed, rank=rank, world_size=world_size, ) - systems = process_systems( - systems_raw, - patterns=dataset_params.get("rglob_patterns"), - ) - return DeepmdDataSystem( - systems=systems, - batch_size=dataset_params["batch_size"], - test_size=1, - type_map=type_map, - trn_all_set=True, - sys_probs=dataset_params.get("sys_probs"), - auto_prob_style=dataset_params.get("auto_prob", "prob_sys_size"), + + return get_data( + dataset_params, 0, type_map, None, lmdb_factory=build_lmdb, seed=seed ) @@ -195,13 +144,16 @@ def get_trainer( is_distributed = dist.is_available() and dist.is_initialized() rank = dist.get_rank() if is_distributed else 0 world_size = dist.get_world_size() if is_distributed else 1 + if data_seed is not None: + seed32 = data_seed % (2**32) + dp_random.seed([rank, seed32] if is_distributed else seed32) def factory( task_config: TrainingTaskConfig, ) -> tuple[DeepmdDataSystem | LmdbDataSystem, Any | None, StatFileSpec]: type_map = list(task_config.model_params["type_map"]) train_data = _build_data_system( - dict(task_config.training_data_params), + task_config.training_data_params, type_map, seed=data_seed, rank=rank, @@ -379,8 +331,8 @@ def update_neighbor_stat( min_nbor_dist: dict[str, float | None] = {} for task_config in iter_training_task_configs(config): type_map = task_config.model_params.get("type_map") - train_data = _get_neighbor_stat_data( - dict(task_config.training_data_params), type_map + train_data = get_neighbor_stat_data( + task_config.training_data_params, type_map ) updated_model_params, task_min_nbor_dist = BaseModel.update_sel( train_data, type_map, dict(task_config.model_params) @@ -528,6 +480,9 @@ def freeze( output: str = "frozen_model", head: str | None = None, lower_kind: str = "nlist", + *, + with_comm: bool = True, + compress: bool = False, ) -> None: """Freeze a pt_expt checkpoint into a .pte/.pt2 exported model. @@ -558,6 +513,14 @@ def freeze( graph lower, so ``lower_kind`` is resolved to ``"graph"`` for it regardless of the requested value -- before the export ABI and the default output suffix are selected. + with_comm : bool, optional + Include a communication artifact for models with cross-rank message + passing, by default True. Disable for single-process inference. + compress : bool, optional + Compress the loaded model before export, by default False. Uses the + default settings of ``dp compress`` and selects the accelerated graph + schema for eligible ``.pt2`` models without compiling an intermediate + artifact. """ import torch @@ -668,12 +631,26 @@ def freeze( ) do_atomic_virial = True + if compress: + from deepmd.pt_expt.entrypoints.compress import ( + compress_model, + ) + + compress_model( + m, + output, + model_def_script=single_model_params, + with_comm=with_comm, + ) + return + model_dict_serialized = m.serialize() deserialize_to_file( output, {"model": model_dict_serialized, "model_def_script": single_model_params}, do_atomic_virial=do_atomic_virial, lower_kind=lower_kind, + with_comm=with_comm, ) log.info("Saved frozen model to %s (lower_kind=%s)", output, lower_kind) @@ -852,7 +829,10 @@ def change_bias( 0.0 if pair_filter is None else float(pair_filter.default), ) model_to_change = model_change_out_bias( - model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode + model_to_change, + sampled_data, + _bias_adjust_mode=bias_adjust_mode, + preset_out_bias={}, ) if input_file.endswith(".pt"): @@ -871,8 +851,19 @@ def change_bias( ) ) model_dict = model_to_change.serialize() + # The updated archive keeps the lower and the per-atom virial of the + # input archive, so that only the bias changes. A compact canonical + # lower serves CUDA devices alone and is selected anew for the host, + # as compression selects it: it stays on a CUDA device, and a host + # without one exports the graph lower. + lower_kind = pte_data["lower_input_kind"] + if lower_kind in ("dpa1_canonical", "dpa4c_canonical"): + lower_kind = "auto" deserialize_to_file( - output_path, {"model": model_dict, "model_def_script": model_params} + output_path, + {"model": model_dict, "model_def_script": model_params}, + do_atomic_virial=pte_data["do_atomic_virial"], + lower_kind=lower_kind, ) log.info(f"Saved model to {output_path}") @@ -942,6 +933,7 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None: output=FLAGS.output, head=FLAGS.head, lower_kind=getattr(FLAGS, "lower_kind", "nlist"), + compress=getattr(FLAGS, "compress", False), ) elif FLAGS.command == "change-bias": change_bias( diff --git a/deepmd/pt_expt/fitting/__init__.py b/deepmd/pt_expt/fitting/__init__.py index 63bcc6c2ce..2f42a57d48 100644 --- a/deepmd/pt_expt/fitting/__init__.py +++ b/deepmd/pt_expt/fitting/__init__.py @@ -17,6 +17,9 @@ from .invar_fitting import ( InvarFitting, ) +from .latent_charge import ( + LatentChargeFitting, +) from .polarizability_fitting import ( PolarFitting, ) @@ -33,6 +36,7 @@ "DipoleFitting", "EnergyFittingNet", "InvarFitting", + "LatentChargeFitting", "PolarFitting", "PropertyFittingNet", "SeZMEnergyFittingNet", diff --git a/deepmd/pt_expt/fitting/latent_charge.py b/deepmd/pt_expt/fitting/latent_charge.py new file mode 100644 index 0000000000..f4f38d2f0a --- /dev/null +++ b/deepmd/pt_expt/fitting/latent_charge.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting as LatentChargeFittingDP, +) +from deepmd.pt_expt.common import ( + torch_module, +) + + +@torch_module +class LatentChargeFitting(LatentChargeFittingDP): + """Per-atom latent charges predicted from the invariant descriptor (pt_expt).""" diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 94c1e1f635..7b658196f0 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -22,6 +22,14 @@ OutputVariableCategory, OutputVariableDef, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, + build_topologies, + check_complete_frames, + check_periodic_cells, + flatten_topologies, + unflatten_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -63,12 +71,14 @@ charge_states_per_frame, single_charge_state, ) +from deepmd.pt_expt.utils.auto_neighbor_list import ( + make_neighbor_list, +) from deepmd.pt_expt.utils.edge_schema import ( edge_schema_from_extended, ) -from deepmd.pt_expt.utils.vesin_neighbor_list import ( - VesinNeighborList, - is_vesin_torch_available, +from deepmd.pt_expt.utils.vesin_search import ( + VesinNeighborSearch, ) from deepmd.utils.charge_state import ( CHARGE_STATE_TABLE_RANGES, @@ -86,6 +96,9 @@ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) + from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + ) # Public output keys emitted by graph-lower forwards, keyed by the @@ -179,25 +192,28 @@ class DeepEval(DeepEvalBackend): nlist_backend : str, default: "auto" Neighbor-list builder for the NLIST/extended lower path (``.pte``, nlist-form ``.pt2``, and dense-routed ``.pt`` checkpoints): - ``"auto"`` / ``"vesin"`` / ``"native"``. Explicit non-default values - are rejected for graph-routed artifacts. + ``"auto"`` / ``"vesin"`` / ``"nv"`` / ``"native"``. Explicit + non-default values are rejected for graph-routed artifacts. neighbor_graph_method : str, default: "auto" Carry-all graph builder for graph-form ``.pt2`` artifacts and graph-routed ``.pt`` checkpoints (``metadata["lower_input_kind"] == "graph"``): ``"auto"`` selects via :func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder` - at each eval call (CUDA: ``nv`` if importable; else ``vesin`` only when - ``nf == 1`` and importable; else ``dense``). Explicit + at each eval call (CUDA: ``vesin`` for one frame if importable, then + ``nv`` if importable, otherwise ``dense``). Explicit ``"dense"`` / ``"ase"`` / ``"cell"`` / ``"vesin"`` / ``"nv"`` choices are preserved. The CPU-only ``cell`` search runs on the host; its graph tensors are transferred to the model device before inference. - A non-default value on any other artifact raises at construction because - the knob would silently do nothing there; use ``nlist_backend`` for the - nlist path instead. All builders emit the same neighbor set, so the - choice is performance-only. Training keeps a separate auto policy + Edge-vector artifacts accept ``"vesin"`` and ``"nv"`` through the same + option. A non-default value on a nlist artifact raises at construction; + use ``nlist_backend`` for that path. All builders emit the same neighbor + set, so the choice is performance-only. Training keeps a separate auto policy (:func:`~deepmd.pt_expt.utils.graph_builder.resolve_neighbor_graph_method`) that never selects ``vesin``. Consolidating the two knobs into a single backend-selection API is deferred to the dense-nlist deprecation. + neighbor_skin : float, default: 0.0 + Vesin candidate-search skin in angstrom, shared by graph, edge-vector, + and extended-coordinate inputs. Zero disables Verlet caching. **kwargs : dict Keyword arguments. """ @@ -211,14 +227,24 @@ def __init__( neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, nlist_backend: str = "auto", neighbor_graph_method: str = "auto", + neighbor_skin: float = 0.0, **kwargs: Any, ) -> None: + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + self.output_def = output_def self.model_path = model_file self.neighbor_list = neighbor_list # Graph-lower builder selection is resolved once after model metadata # identifies the lower ABI. self._neighbor_graph_method = neighbor_graph_method + self._neighbor_skin = neighbor_skin + # CUDA destination grouping belongs to the graph CSR construction. + self._vesin_search = VesinNeighborSearch( + skin=neighbor_skin, sorted=DEVICE.type != "cuda" + ) self._is_pt2 = model_file.endswith(".pt2") # Only a compressed ``.pt2`` folds its charge state into constants; a # model that reads the condition as an ordinary input needs no rebuild. @@ -241,22 +267,34 @@ def __init__( # above (_load_pt2 / _load_pte / _load_pt), so this is the one # place that sees every model load regardless of archive kind. _warn_legacy_edge_vec(self.metadata) + if self.metadata.get("lower_input_kind") in ( + "dpa1_canonical", + "dpa4c_canonical", + ) and ( + self.metadata.get("canonical_index_dtype") != "auto" + or self.metadata.get("canonical_index_storage") != "byte_words" + ): + raise ValueError( + "Compact canonical graph artifacts require automatic-width byte " + "indices; re-freeze the model with the current DeePMD-kit version." + ) - # ``neighbor_graph_method`` is consumed only by graph-lower evaluation. - # Fail fast instead of silently ignoring it on nlist-form artifacts, - # where the corresponding builder knob is ``nlist_backend``. + # ``neighbor_graph_method`` selects searches for graph and edge-vector + # inputs. Extended-coordinate inputs use ``nlist_backend`` instead. if neighbor_graph_method != "auto" and getattr(self, "metadata", {}).get( "lower_input_kind" ) not in ( "graph", "dpa1_canonical", "dpa4c_canonical", + "edge_vec", ): raise ValueError( f"neighbor_graph_method={neighbor_graph_method!r} only applies to " - "graph-routed artifacts (lower_input_kind == 'graph'); this " - f"model is not graph-form. Use nlist_backend to select the " - "neighbor-list builder for the nlist path." + "graph or edge-vector artifacts (lower_input_kind: 'graph', " + "'dpa1_canonical', 'dpa4c_canonical', or 'edge_vec'). " + f"This model uses {self.metadata.get('lower_input_kind')!r}; " + "use nlist_backend to select its neighbor-list builder." ) self._setup_neighbor_backend(nlist_backend) @@ -303,18 +341,25 @@ def _resolve_neighbor_graph_method(method: str, nf: int | None = None) -> str: def _setup_neighbor_backend(self, nlist_backend: str) -> None: """Resolve the graph or neighbor-list construction strategy. - ``"native"`` uses the dense all-pairs builder; ``"vesin"`` forces the - O(N) ``vesin.torch`` cell list (raising if it is unavailable or the - model/inputs are unsupported); ``"auto"`` uses vesin when applicable and - silently falls back to the native builder otherwise. Graph-routed - artifacts use ``neighbor_graph_method`` instead and reject nlist-specific - controls rather than silently changing or ignoring the requested lower. + ``"native"`` retains unbuffered construction and ignores + ``neighbor_skin``: edge-vector artifacts use the batch-aware Vesin/NV + policy, with dense carry-all edges only when both optional libraries + are unavailable; nlist artifacts retain native extended coordinates. + ``"vesin"`` and ``"nv"`` select the corresponding search, raising for + unavailable or unsupported inputs. ``"auto"`` selects by frame count and + available libraries, falling back to the native builder when none applies. + Graph-routed artifacts use ``neighbor_graph_method`` instead and reject + nlist-specific controls rather than silently changing or ignoring the requested lower. Results are unchanged either way; only the neighbor-search cost differs. """ - if nlist_backend not in ("auto", "vesin", "native"): + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + if nlist_backend not in ("auto", "vesin", "nv", "native"): raise ValueError( f"Unknown nlist_backend '{nlist_backend}'; " - "expected 'auto', 'vesin', or 'native'." + "expected 'auto', 'vesin', 'nv', or 'native'." ) if self.metadata.get("lower_input_kind") in ( "graph", @@ -339,40 +384,57 @@ def _setup_neighbor_backend(self, nlist_backend: str) -> None: self._neighbor_graph_method = self._resolve_neighbor_graph_method( self._neighbor_graph_method ) - self._use_vesin = False self._nlist_builder = None return + if ( + self.metadata.get("lower_input_kind") == "edge_vec" + and self._neighbor_graph_method != "auto" + ): + if self._neighbor_graph_method not in ("vesin", "nv"): + raise ValueError( + "Edge-vector artifacts support neighbor_graph_method " + "'auto', 'vesin', or 'nv'; " + f"got {self._neighbor_graph_method!r}." + ) + if nlist_backend != "auto": + raise ValueError( + "Select only one of neighbor_graph_method and nlist_backend." + ) + nlist_backend = self._neighbor_graph_method is_spin = bool(getattr(self, "_is_spin", False)) - # Graph-routed artifacts returned above. Any remaining spin model uses - # the virtual-atom nlist scheme, which vesin does not support. + # Native-spin edge artifacts share the energy neighbor strategy. + # Virtual-atom spin expansion retains the native nlist path. ase_provided = self.neighbor_list is not None - # reason vesin cannot be used (None means it can) - unsupported = "spin models" if is_spin else None - if nlist_backend == "native": - self._use_vesin = False - elif nlist_backend == "vesin": - if not is_vesin_torch_available(): - raise ImportError( - "nlist_backend='vesin' was requested but 'vesin.torch' is " - "not installed. Install it (`pip install vesin[torch]`) or " - "use nlist_backend='native' (or 'auto')." - ) + unsupported = ( + "spin models" + if is_spin and self.metadata.get("lower_input_kind") != "edge_vec" + else None + ) + if nlist_backend in ("vesin", "nv"): if unsupported is not None: raise ValueError( - f"nlist_backend='vesin' is not supported for {unsupported}; " + f"nlist_backend={nlist_backend!r} is not supported for {unsupported}; " "use nlist_backend='native' (or 'auto')." ) if ase_provided: raise ValueError( - "nlist_backend='vesin' conflicts with an explicitly " + f"nlist_backend={nlist_backend!r} conflicts with an explicitly " "supplied ASE neighbor_list; pass only one." ) - self._use_vesin = True - else: # auto: use vesin when possible, otherwise fall back silently - self._use_vesin = ( - is_vesin_torch_available() and unsupported is None and not ase_provided + self._nlist_builder = ( + make_neighbor_list( + nlist_backend, + DEVICE, + self._neighbor_skin, + return_mode=( + "edges" + if self.metadata.get("lower_input_kind") == "edge_vec" + else "extended" + ), ) - self._nlist_builder = VesinNeighborList() if self._use_vesin else None + if unsupported is None and not ase_provided + else None + ) def _init_from_model_json(self, model_json_str: str) -> None: """Deserialize model.json and derive model API from the dpmodel instance.""" @@ -473,7 +535,10 @@ def _load_pte(self, model_file: str) -> None: "model_def_script.json": "", "metadata.json": "", } - exported = torch.export.load(model_file, extra_files=extra_files) + # Archive loading stages raw storages on CPU before restoring their + # recorded devices, including when the caller selects another default. + with torch.device("cpu"): + exported = torch.export.load(model_file, extra_files=extra_files) self.exported_module = exported.module() mds = extra_files["model_def_script.json"] self._model_def_script = json.loads(mds) if mds else {} @@ -664,6 +729,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: ) use_graph_lower = model_uses_graph_lower(model) + topology_names = model.atomic_model.required_topologies() self._dpmodel = model self._is_spin = ( model_params.get("type") == "spin_ener" or "spin" in model_params @@ -718,6 +784,9 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "is_spin": self._is_spin, "lower_input_kind": "graph" if use_graph_lower else "nlist", } + if topology_names: + self.metadata["topologies"] = list(topology_names) + self.metadata["topology_params"] = model.atomic_model.topology_params() if use_graph_lower: from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, @@ -735,7 +804,8 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: # ``communicate_extended_output``, while the graph runner emits local # public keys consumed directly by ``_eval_model_graph``. # - # Graph: (..., source_row_ptr, fparam, aparam, charge_spin) + # Graph: (..., source_row_ptr, fparam, aparam, charge_spin, + # *topology inputs) # Graph spin: (..., source_row_ptr, spin, fparam, aparam, charge_spin) # Nlist: (ext_coord, ext_atype, nlist, mapping, fparam, aparam) # Nlist spin: (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) @@ -805,6 +875,7 @@ def _eager_runner_graph( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None = None, + *topology_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: model_ret = model.forward_common_lower_graph( atype, @@ -822,6 +893,7 @@ def _eager_runner_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + topologies=unflatten_topologies(topology_names, topology_vals), ) return _translate_energy_keys( model_ret, @@ -1639,28 +1711,6 @@ def _build_nlist_ase_single( return extended_coord, extended_atype, nlist, mapping - @staticmethod - def _build_edge_inputs_from_nlist( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Convert a padded neighbor list into the compact-edge schema.""" - nloc = nlist.shape[1] - schema = edge_schema_from_extended( - extended_coord, - extended_atype[:, :nloc], - nlist, - mapping, - ) - return ( - schema.edge_index, - schema.edge_vec, - schema.edge_scatter_index, - schema.edge_mask, - ) - def _prepare_nlist_inputs( self, coords: np.ndarray, @@ -1719,47 +1769,9 @@ def _prepare_nlist_inputs( atype_t, ) - if fparam is not None: - fparam_t = torch.tensor( - fparam.reshape(nframes, self.get_dim_fparam()), - dtype=torch.float64, - device=DEVICE, - ) - elif self.get_dim_fparam() > 0: - # Exported models (.pt2/.pte) are compiled with fparam as a - # required input. Fill with default values from metadata. - default_fp = self.metadata.get("default_fparam") - if default_fp is not None: - fparam_t = ( - torch.tensor(default_fp, dtype=torch.float64, device=DEVICE) - .unsqueeze(0) - .expand(nframes, -1) - .contiguous() - ) - else: - raise ValueError( - f"fparam is required for this model (dim_fparam={self.get_dim_fparam()}) " - "but was not provided, and no default_fparam is stored in the model." - ) - else: - fparam_t = None - - if aparam is not None: - aparam_t = torch.tensor( - aparam.reshape(nframes, natoms, self.get_dim_aparam()), - dtype=torch.float64, - device=DEVICE, - ) - elif self.get_dim_aparam() > 0: - # Exported models (.pt2/.pte) are compiled with aparam as a - # required positional input. Unlike fparam, there is no default. - raise ValueError( - f"aparam is required for this model (dim_aparam={self.get_dim_aparam()}) " - "but was not provided." - ) - else: - aparam_t = None - + fparam_t, aparam_t = self._prepare_optional_lower_inputs( + fparam, aparam, nframes, natoms, DEVICE + ) charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) return ( @@ -1857,21 +1869,19 @@ def _prepare_inputs( coords, cells, atom_types, fparam, aparam, charge_spin ) if self.metadata.get("lower_input_kind") == "edge_vec": - edge_index_t, edge_vec_t, edge_scatter_t, edge_mask_t = ( - self._build_edge_inputs_from_nlist( - ext_coord_t, - ext_atype_t, - nlist_t, - mapping_t, - ) - ) - model_inputs = ( + edge_schema = edge_schema_from_extended( ext_coord_t, ext_atype_t[:, :natoms], - edge_index_t, - edge_vec_t, - edge_scatter_t, - edge_mask_t, + nlist_t, + mapping_t, + ) + model_inputs = ( + edge_schema.coord, + edge_schema.atype, + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_scatter_index, + edge_schema.edge_mask, fparam_t, aparam_t, charge_spin_t, @@ -1904,6 +1914,8 @@ def _prepare_optional_lower_inputs( device=device, ) elif self.get_dim_fparam() > 0: + # Compiled lower interfaces require a tensor for frame parameters; + # omitted values use the defaults stored in the model metadata. default_fp = self.metadata.get("default_fparam") if default_fp is None: raise ValueError( @@ -1926,6 +1938,7 @@ def _prepare_optional_lower_inputs( device=device, ) elif self.get_dim_aparam() > 0: + # Atomic parameters have no model-level default. raise ValueError( f"aparam is required for this model (dim_aparam={self.get_dim_aparam()}) " "but was not provided." @@ -2022,130 +2035,29 @@ def _eval_model_spin( request_defs, charge_spin=charge_spin, ) - nframes = coords.shape[0] - if len(atom_types.shape) == 1: - natoms = len(atom_types) - atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) - else: - natoms = len(atom_types[0]) - from deepmd.pt_expt.utils.env import ( DEVICE, ) - coord_input = coords.reshape(nframes, natoms, 3) - if self.neighbor_list is not None: - extended_coord, extended_atype, nlist, mapping = self._build_nlist_ase( - coord_input, - cells, - atom_types, - ) - ext_coord_t = torch.tensor( - extended_coord, dtype=torch.float64, device=DEVICE - ) - ext_atype_t = torch.tensor(extended_atype, dtype=torch.int64, device=DEVICE) - nlist_t = torch.tensor(nlist, dtype=torch.int64, device=DEVICE) - mapping_t = torch.tensor(mapping, dtype=torch.int64, device=DEVICE) - else: - coord_t = torch.tensor(coord_input, dtype=torch.float64, device=DEVICE) - atype_t = torch.tensor(atom_types, dtype=torch.int64, device=DEVICE) - cells_t = ( - torch.tensor(cells, dtype=torch.float64, device=DEVICE) - if cells is not None - else None - ) - ext_coord_t, ext_atype_t, nlist_t, mapping_t = self._build_nlist_native( - coord_t, - cells_t, - atype_t, - ) - - # Extend spin to ghost atoms using mapping + model_inputs, mapping_t, nframes, natoms = self._prepare_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin + ) spin_t = torch.tensor( spins.reshape(nframes, natoms, 3), dtype=torch.float64, device=DEVICE ) - batch_idx = ( - torch.arange(nframes, dtype=torch.long, device=DEVICE) - .unsqueeze(1) - .expand_as(mapping_t) - ) - ext_spin_t = spin_t[batch_idx, mapping_t] - - if fparam is not None: - fparam_t = torch.tensor( - fparam.reshape(nframes, self.get_dim_fparam()), - dtype=torch.float64, - device=DEVICE, - ) - elif self.get_dim_fparam() > 0: - # Exported models (.pt2/.pte) are compiled with fparam as a - # required input. Fill with default values from metadata. - default_fp = self.metadata.get("default_fparam") - if default_fp is not None: - fparam_t = ( - torch.tensor(default_fp, dtype=torch.float64, device=DEVICE) - .unsqueeze(0) - .expand(nframes, -1) - .contiguous() - ) - else: - raise ValueError( - f"fparam is required for this model (dim_fparam={self.get_dim_fparam()}) " - "but was not provided, and no default_fparam is stored in the model." - ) - else: - fparam_t = None - - if aparam is not None: - aparam_t = torch.tensor( - aparam.reshape(nframes, natoms, self.get_dim_aparam()), - dtype=torch.float64, - device=DEVICE, - ) - elif self.get_dim_aparam() > 0: - raise ValueError( - f"aparam is required for this model (dim_aparam={self.get_dim_aparam()}) " - "but was not provided." - ) - else: - aparam_t = None - - charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) - - # Build the lower inputs for the model's spin ABI. The native scheme - # shares the energy edge contract and feeds the owned-atom spins (the - # descriptor only needs local spins; ghost neighbours resolve to their - # local owners). The deepspin scheme keeps the extended nlist contract. + # Native spin belongs to the local owners of the edge schema. The + # virtual-atom scheme extends it to ghost sites through the same mapping + # used to fold the model's coordinate derivatives back to local atoms. if self.metadata.get("lower_input_kind") == "edge_vec": - edge_schema = edge_schema_from_extended( - ext_coord_t, - ext_atype_t[:, :natoms], - nlist_t, - mapping_t, - ) - model_inputs = ( - edge_schema.coord, - edge_schema.atype, - edge_schema.edge_index, - edge_schema.edge_vec, - edge_schema.edge_scatter_index, - edge_schema.edge_mask, - spin_t, - fparam_t, - aparam_t, - charge_spin_t, - ) + model_inputs = (*model_inputs[:6], spin_t, *model_inputs[6:]) else: - model_inputs = ( - ext_coord_t, - ext_atype_t, - ext_spin_t, - nlist_t, - mapping_t, - fparam_t, - aparam_t, - charge_spin_t, + batch_idx = ( + torch.arange(nframes, dtype=torch.long, device=DEVICE) + .unsqueeze(1) + .expand_as(mapping_t) ) + ext_spin_t = spin_t[batch_idx, mapping_t] + model_inputs = (*model_inputs[:2], ext_spin_t, *model_inputs[2:]) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: @@ -2200,158 +2112,16 @@ def _eval_model_graph_spin( request_defs: list[OutputVariableDef], charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: - """Evaluate a graph-form native-spin ``.pt2`` (``lower_input_kind == - "graph"`` and ``is_spin``). - - Mirrors :meth:`_eval_model_graph`'s carry-all - :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph` - construction (SAME builder, SAME positional ABI up through - ``source_row_ptr``), then inserts the owned-atom ``spin`` tensor - ``(N, 3)`` at positional index 10 of - ``NativeSpinEnergyModel.forward_lower_graph_exportable`` -- the node - axis IS the owned-local-atom axis for single-rank eval (no ghost - nodes), so ``spin`` needs no extension/mapping, unlike the dense - spin path's ``ext_spin_t``. ``charge_spin`` rides the conditional - slot-13 tail (combined native-spin + charge-spin FiLM models; the - slot is dropped from the exported signature otherwise). The forward - returns LOCAL public keys directly (``atom_energy``, - ``energy``, ``force``, ``force_mag``, ``virial``, ``atom_virial``), - so results are reshaped without ``communicate_extended_output``, - same as the non-spin graph path. - """ - from deepmd.pt_expt.utils.env import ( - DEVICE, - ) + """Evaluate native-spin graph or compact canonical inputs. - nframes = coords.shape[0] - if len(atom_types.shape) == 1: - natoms = len(atom_types) - atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) - else: - natoms = len(atom_types[0]) - - coord_input = coords.reshape(nframes, natoms, 3) - box_input = cells.reshape(nframes, 9) if cells is not None else None - graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE) - - atype_t = torch.tensor( - np.asarray(atom_types).reshape(-1), dtype=torch.int64, device=DEVICE - ) - n_node_t = torch.as_tensor(graph.n_node, dtype=torch.int64, device=DEVICE) - edge_index_t = torch.as_tensor( - graph.edge_index, dtype=torch.int64, device=DEVICE - ) - edge_dtype = ( - torch.float32 - if self.metadata.get("graph_edge_dtype") == "float32" - else torch.float64 - ) - edge_vec_t = torch.as_tensor( - graph.edge_vec, - dtype=edge_dtype, - device=DEVICE, - ) - edge_mask_t = torch.as_tensor(graph.edge_mask, dtype=torch.bool, device=DEVICE) - destination_order_t = torch.as_tensor( - graph.destination_order, - device=DEVICE, - ) - destination_row_ptr_t = torch.as_tensor( - graph.destination_row_ptr, - dtype=torch.int64, - device=DEVICE, - ) - source_order_t = torch.as_tensor( - graph.source_order, - device=DEVICE, - ) - source_row_ptr_t = torch.as_tensor( - graph.source_row_ptr, - dtype=torch.int64, - device=DEVICE, - ) - - spin_t = torch.tensor( - np.asarray(spins).reshape(nframes * natoms, 3), - dtype=torch.float64, - device=DEVICE, - ) - - fparam_t, aparam_t = self._prepare_optional_lower_inputs( - fparam, aparam, nframes, natoms, DEVICE + Spin belongs to the local node axis, so no ghost extension is needed. + The shared input preparation places spin before the generic graph's + optional conditioning inputs, or last in the compact canonical ABI. + Both forwards return local public outputs directly. + """ + model_inputs, nframes, natoms = self._prepare_graph_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin, spins=spins ) - if aparam_t is not None: - # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) -- - # the same axis as ``atype``/``spin`` (mirrors _eval_model_graph). - aparam_t = aparam_t.reshape(nframes * natoms, -1) - - if self.metadata.get("lower_input_kind") == "dpa4c_canonical": - # The compact canonical ABI has no conditional tail, so the moment - # is the last slot rather than the eleventh, and any conditioning - # input the artifact cannot receive is an error rather than a - # silently dropped argument. - if ( - self.get_dim_fparam() > 0 - or self.get_dim_aparam() > 0 - or int(self.metadata.get("dim_chg_spin", 0) or 0) > 0 - ): - raise NotImplementedError( - "compact canonical artifacts carry no fparam/aparam/" - "charge_spin inputs; a model requiring them must not be " - "frozen with a canonical lower kind." - ) - # A compressed descriptor still serves a condition, through the - # constants its fold rewrites rather than through an argument. - self._apply_charge_state(charge_spin) - from deepmd.dpmodel.utils.neighbor_graph import ( - NeighborGraph, - ) - from deepmd.pt_expt.utils.canonical_graph import ( - canonical_graph_from_neighbor_graph, - ) - - compact = canonical_graph_from_neighbor_graph( - NeighborGraph( - n_node=n_node_t, - edge_index=edge_index_t, - edge_vec=edge_vec_t, - edge_mask=edge_mask_t, - n_local=n_node_t, - destination_order=destination_order_t, - destination_row_ptr=destination_row_ptr_t, - source_order=source_order_t, - source_row_ptr=source_row_ptr_t, - destination_sorted=bool(graph.destination_sorted), - ) - ) - model_inputs = ( - atype_t, - compact.n_node, - compact.n_local, - compact.source, - compact.edge_vec, - compact.destination_row_ptr, - compact.source_row_ptr, - compact.source_order, - spin_t.to(torch.float32), - ) - else: - model_inputs = ( - atype_t, - n_node_t, - n_node_t, - edge_index_t, - edge_vec_t, - edge_mask_t, - destination_order_t, - destination_row_ptr_t, - source_order_t, - source_row_ptr_t, - spin_t, - fparam_t, - aparam_t, - self._make_charge_spin_input(nframes, charge_spin), - ) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: @@ -2380,85 +2150,117 @@ def _eval_model_graph( request_defs: list[OutputVariableDef], charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: - """Evaluate a graph-lower model (``lower_input_kind == "graph"``). - - Builds a carry-all :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph` - from the eval system at its exact (tight) edge count and feeds the - positional schema - ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, - destination_order, destination_row_ptr, source_order, source_row_ptr, - fparam, aparam, charge_spin)`` to the graph forward. Exported AOTI - artifacts use a dynamic edge axis, so no ``edge_capacity`` padding is - needed. The ``graph_edge_dtype`` metadata selects float32 geometry for - compressed DPA1 and float64 for generic graph descriptors. The forward - returns local public keys directly, so results are reshaped without - ``communicate_extended_output``. + """Evaluate generic graph or compact canonical inputs. + + Exported artifacts accept a dynamic edge axis, and the inputs of a + model with a long-range module, the fields of every topology that the + ``topologies`` metadata names, follow the graph inputs. Input + preparation uses the artifact's geometry dtype and layout; both + forwards return local public outputs without extended-atom + communication. """ + model_inputs, nframes, natoms = self._prepare_graph_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin + ) + if self._is_pt2: + model_ret = self._pt2_runner(*model_inputs) + else: + model_ret = self.exported_module(*model_inputs) + + results = [] + for odef in request_defs: + shape = self._get_output_shape(odef, nframes, natoms) + gkey = _GRAPH_CATEGORY_TO_KEY.get(odef.category) + val = model_ret.get(gkey) if gkey is not None else None + if val is not None: + results.append(val.detach().cpu().numpy().reshape(shape)) + else: + results.append( + np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION) + ) + return tuple(results) + + def _prepare_graph_inputs( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None, + aparam: np.ndarray | None, + charge_spin: np.ndarray | None, + *, + spins: np.ndarray | None = None, + ) -> tuple[tuple[torch.Tensor | None, ...], int, int]: + """Pack local graph inputs according to the exported artifact's ABI.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + canonical_graph_from_neighbor_graph, + canonical_index_bytes, + ) from deepmd.pt_expt.utils.env import ( DEVICE, ) nframes = coords.shape[0] - if len(atom_types.shape) == 1: + if atom_types.ndim == 1: natoms = len(atom_types) - atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) + atom_types = np.tile(atom_types, (nframes, 1)) else: - natoms = len(atom_types[0]) - - coord_input = coords.reshape(nframes, natoms, 3) - box_input = cells.reshape(nframes, 9) if cells is not None else None - # Build the carry-all graph at its exact edge count; the exported edge - # axis is dynamic. - graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE) - - atype_t = torch.tensor( - np.asarray(atom_types).reshape(-1), dtype=torch.int64, device=DEVICE + natoms = atom_types.shape[1] + graph = self._build_eval_graph( + coords.reshape(nframes, natoms, 3), + atom_types, + cells.reshape(nframes, 9) if cells is not None else None, + DEVICE, ) - # graph fields may be numpy (dense/ase) or torch, possibly on CUDA - # (vesin/nv) -- torch.as_tensor handles both and moves to DEVICE. - n_node_t = torch.as_tensor(graph.n_node, dtype=torch.int64, device=DEVICE) - edge_index_t = torch.as_tensor( - graph.edge_index, dtype=torch.int64, device=DEVICE + atype_t = torch.as_tensor( + atom_types.reshape(-1), dtype=torch.int64, device=DEVICE + ) + canonical = self.metadata.get("lower_input_kind") in ( + "dpa1_canonical", + "dpa4c_canonical", ) edge_dtype = ( torch.float32 if self.metadata.get("graph_edge_dtype") == "float32" else torch.float64 ) - edge_vec_t = torch.as_tensor( - graph.edge_vec, - dtype=edge_dtype, - device=DEVICE, - ) - edge_mask_t = torch.as_tensor(graph.edge_mask, dtype=torch.bool, device=DEVICE) - destination_order_t = torch.as_tensor( - graph.destination_order, - device=DEVICE, - ) - destination_row_ptr_t = torch.as_tensor( - graph.destination_row_ptr, - dtype=torch.int64, - device=DEVICE, - ) - source_order_t = torch.as_tensor( - graph.source_order, - device=DEVICE, - ) - source_row_ptr_t = torch.as_tensor( - graph.source_row_ptr, - dtype=torch.int64, - device=DEVICE, - ) - - if self.metadata.get("lower_input_kind") in ( - "dpa1_canonical", - "dpa4c_canonical", - ): - # The canonical ABI has NO fparam/aparam/charge_spin slots; the - # export gate (fitting_eligible) rejects such models today, so - # this is unreachable -- assert it loudly so a future loosening - # of the eligibility fails here instead of silently dropping - # conditioning inputs at inference. + if not isinstance(graph, CanonicalGraph): + # CPU builders return NumPy arrays; device builders return tensors. + # All graph fields cross the model-device boundary together. + n_node = torch.as_tensor(graph.n_node, dtype=torch.int64, device=DEVICE) + graph = NeighborGraph( + n_node=n_node, + n_local=n_node, + edge_index=torch.as_tensor( + graph.edge_index, dtype=torch.int64, device=DEVICE + ), + edge_vec=torch.as_tensor( + graph.edge_vec, dtype=edge_dtype, device=DEVICE + ), + edge_mask=torch.as_tensor( + graph.edge_mask, dtype=torch.bool, device=DEVICE + ), + destination_order=torch.as_tensor( + graph.destination_order, device=DEVICE + ), + destination_row_ptr=torch.as_tensor( + graph.destination_row_ptr, dtype=torch.int64, device=DEVICE + ), + source_order=torch.as_tensor(graph.source_order, device=DEVICE), + source_row_ptr=torch.as_tensor( + graph.source_row_ptr, dtype=torch.int64, device=DEVICE + ), + destination_sorted=bool(graph.destination_sorted), + ) + if canonical: + graph = canonical_graph_from_neighbor_graph(graph) + if canonical: + # The compact ABI has no conditional input slots. Compressed charge + # conditioning instead updates the constants through its fold. if ( self.get_dim_fparam() > 0 or self.get_dim_aparam() > 0 @@ -2467,84 +2269,132 @@ def _eval_model_graph( raise NotImplementedError( "compact canonical artifacts carry no fparam/aparam/" "charge_spin inputs; a model requiring them must not be " - "frozen with a canonical lower kind (the export eligibility " - "gate should have rejected it)." + "frozen with a canonical lower kind." ) - # A compressed descriptor still serves a condition, through the - # constants its fold rewrites rather than through an argument. self._apply_charge_state(charge_spin) - from deepmd.dpmodel.utils.neighbor_graph import ( - NeighborGraph, - ) - from deepmd.pt_expt.utils.canonical_graph import ( - canonical_graph_from_neighbor_graph, - ) - - generic_graph = NeighborGraph( - n_node=n_node_t, - edge_index=edge_index_t, - edge_vec=edge_vec_t, - edge_mask=edge_mask_t, - n_local=n_node_t, - destination_order=destination_order_t, - destination_row_ptr=destination_row_ptr_t, - source_order=source_order_t, - source_row_ptr=source_row_ptr_t, - destination_sorted=bool(graph.destination_sorted), - ) - compact = canonical_graph_from_neighbor_graph(generic_graph) model_inputs = ( atype_t, - compact.n_node, - compact.n_local, - compact.source, - compact.edge_vec, - compact.destination_row_ptr, - compact.source_row_ptr, - compact.source_order, + graph.n_node, + graph.n_local, + canonical_index_bytes(graph.source), + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, ) else: + model_inputs = ( + atype_t, + graph.n_node, + graph.n_local, + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + ) + if spins is not None: + spin_t = torch.as_tensor( + spins.reshape(nframes * natoms, 3), + dtype=torch.float32 if canonical else torch.float64, + device=DEVICE, + ) + model_inputs = (*model_inputs, spin_t) + if not canonical: fparam_t, aparam_t = self._prepare_optional_lower_inputs( fparam, aparam, nframes, natoms, DEVICE ) if aparam_t is not None: - # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) - # -- the same axis as ``atype`` (the shared helper above - # returns the dense rectangular (nf, natoms, nda) layout). + # Graph inputs flatten frame and atom axes into one node axis. aparam_t = aparam_t.reshape(nframes * natoms, -1) - charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) model_inputs = ( - atype_t, - n_node_t, - n_node_t, - edge_index_t, - edge_vec_t, - edge_mask_t, - destination_order_t, - destination_row_ptr_t, - source_order_t, - source_row_ptr_t, + *model_inputs, fparam_t, aparam_t, - charge_spin_t, + self._make_charge_spin_input(nframes, charge_spin), ) - if self._is_pt2: - model_ret = self._pt2_runner(*model_inputs) - else: - model_ret = self.exported_module(*model_inputs) + topology_names = self.metadata.get("topologies", ()) + if topology_names: + box_input = cells.reshape(nframes, 9) if cells is not None else None + model_inputs = ( + *model_inputs, + *self._long_range_inputs( + topology_names, + coords.reshape(nframes * natoms, 3), + box_input, + atype_t, + graph.n_node, + canonical, + edge_dtype, + DEVICE, + ), + ) + return model_inputs, nframes, natoms - results = [] - for odef in request_defs: - shape = self._get_output_shape(odef, nframes, natoms) - gkey = _GRAPH_CATEGORY_TO_KEY.get(odef.category) - val = model_ret.get(gkey) if gkey is not None else None - if val is not None: - results.append(val.detach().cpu().numpy().reshape(shape)) - else: - results.append( - np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION) + def _long_range_inputs( + self, + topology_names: tuple[str, ...], + coord: np.ndarray, + box_input: np.ndarray | None, + atype: torch.Tensor, + n_node: torch.Tensor, + canonical: bool, + edge_dtype: torch.dtype, + device: "torch.device", + ) -> tuple[torch.Tensor | None, ...]: + """Inputs that the lower of a model with a long-range module reads after the graph. + + The compact lower reads the positions and the cells of the frames, and + the zero communicator of frames held whole; a module that serves no + periodic system takes no cell, and the zero cell that marks a + non-periodic frame is never one a caller passes. The graph lower reads + the fields of the topologies, built with the ``topology_params`` of + the archive, or those of the live model of a checkpoint. + """ + nframes = n_node.shape[0] + if canonical: + if PME_MESH not in topology_names: + check_complete_frames(box_input) + elif box_input is not None: + check_periodic_cells(box_input) + box_t = ( + torch.zeros(nframes, 9, dtype=torch.float64, device="cpu") + if box_input is None + else torch.as_tensor( + np.asarray(box_input, dtype=np.float64).reshape(nframes, 9) ) - return tuple(results) + ) + return ( + torch.as_tensor(coord, dtype=torch.float64, device=device), + torch.zeros(1, dtype=torch.int64, device="cpu"), + box_t, + ) + # Padding slots of a mixed-type batch are no atoms and join no pair. + # The cell is torch like the coordinates: the topologies it + # namespaces feed the torch graph lower directly (a numpy cell would + # produce numpy leaves). + box_t = ( + torch.as_tensor( + np.asarray(box_input).reshape(nframes, 3, 3), + dtype=edge_dtype, + device=device, + ) + if box_input is not None + else None + ) + topologies = build_topologies( + topology_names, + torch.as_tensor(coord, dtype=edge_dtype, device=device), + n_node, + box_t, + node_mask=atype >= 0, + atype=atype, + pair_excl=self._model_pair_excl(), + topology_params=self.metadata.get("topology_params"), + ) + return flatten_topologies(topologies) def _build_eval_graph( self, @@ -2552,8 +2402,8 @@ def _build_eval_graph( atom_types: np.ndarray, box_input: np.ndarray | None, device: "torch.device", - ) -> "NeighborGraph": - """Build the carry-all NeighborGraph for graph-lower inference. + ) -> "NeighborGraph | CanonicalGraph": + """Build the local graph for the artifact's input layout. Dispatches on ``self._neighbor_graph_method``: ``auto`` is resolved call-time via @@ -2565,7 +2415,8 @@ def _build_eval_graph( (carry-all, sel-free), so the selection is a pure performance choice and results are unchanged. The result is canonicalized to the destination-major graph-form ``.pt2`` ABI after construction; the - caller transfers its fields to the model device. + caller transfers its fields to the model device. Vesin emits the + compact canonical layout directly when no pair exclusions are needed. """ method = self._neighbor_graph_method if method == "auto": @@ -2672,9 +2523,23 @@ def _build_eval_graph( ) if method == "vesin": from deepmd.pt_expt.utils.vesin_graph_builder import ( + build_canonical_graph_vesin, build_neighbor_graph_vesin, ) + if pair_excl is None and self.metadata.get("lower_input_kind") in ( + "dpa1_canonical", + "dpa4c_canonical", + ): + return build_canonical_graph_vesin( + cc, + bb.reshape(-1, 3, 3) if bb is not None else None, + self._rcut, + self._vesin_search, + node_mask=aa >= 0 + if (np.asarray(atom_types) < 0).any() + else None, + ) return build_neighbor_graph_vesin( cc, aa, @@ -2682,6 +2547,7 @@ def _build_eval_graph( self._rcut, canonicalize=True, pair_excl=pair_excl, + search=self._vesin_search, ) from deepmd.pt_expt.utils.nv_graph_builder import ( build_neighbor_graph_nv, @@ -2703,10 +2569,12 @@ def _build_eval_graph( def _model_pair_excl(self) -> "PairExcludeMask | None": """Model-level ``pair_exclude_types`` as a ``PairExcludeMask`` (or None). - Applied at graph BUILD time (decision #18), NOT inside the exported - ``.pt2`` lower. Reads the excluded pairs from the loaded dpmodel (if any) - or the ``pair_exclude_types`` field in ``metadata.json``, and returns a - FRESH numpy-backed mask. + Graph and nlist feeders apply these exclusions at build time: their + exported lowers consume pre-excluded inputs. PT ``edge_vec`` artifacts + instead retain descriptor-owned exclusion inside the compiled math. + Reads the excluded pairs from the loaded dpmodel (if any) or the + ``pair_exclude_types`` field in ``metadata.json``, and returns a fresh + NumPy-backed mask. A numpy ``type_mask`` converts cleanly onto whichever namespace/device the builder's ``atype`` uses (dense/ase pass numpy; cell/vesin/nv pass torch). diff --git a/deepmd/pt_expt/kernels/autotune.py b/deepmd/pt_expt/kernels/autotune.py index b635590e01..4e377f9ae4 100644 --- a/deepmd/pt_expt/kernels/autotune.py +++ b/deepmd/pt_expt/kernels/autotune.py @@ -1,17 +1,20 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Freeze-time Triton launch-configuration autotuning. +"""Freeze-time autotuning of custom-kernel launch choices. Custom-kernel families register a tuner through :func:`register_autotuner`; a model-freezing entry point calls :func:`run_autotune` once, before tracing, so -the exported artifact bakes launch configurations tuned for the exact -deployment shapes and GPU. A tuner sweeps only the shape keys its kernels need -that are not already covered by the built-in tables, keeping repeat freezes -cheap. - -Autotuning is a no-op below ``DP_TRITON_INFER`` level 2 (level-1 kernels run a -single shape-independent configuration) and off CUDA -- Triton launch tables -are GPU-specific and AOTInductor artifacts are not portable across GPU models, -so the tuning target is always the local deployment GPU. +the exported artifact bakes launch choices tuned for the exact deployment +shapes and GPU: the Triton launch configurations of the DPA1 kernels, and the +fitting routes of the fused energy-force operator of compressed DPA4C. A tuner +measures only the keys its kernels need that are not already covered, keeping +repeat freezes cheap. + +Each tuner runs from the ``DP_TRITON_INFER`` level it registers with (level-1 +Triton kernels run a single shape-independent configuration, so the Triton +tuners start at level 2), for the lower kinds whose programs call its kernels, +and only on CUDA: the tables are GPU-specific and AOTInductor artifacts are not +portable across GPU models, so the tuning target is always the local +deployment GPU. Adding a new custom-kernel family is one registration call at import time; the freeze paths need no change. @@ -34,34 +37,51 @@ log = logging.getLogger(__name__) -# Family name -> tuner(model, level, device). Keying by name makes repeated -# registration (e.g. module re-import) idempotent. +# Family name -> (tuner(model, level, device), lowest level, lower kinds or +# None for all). Keying by name makes repeated registration (e.g. module +# re-import) idempotent. Autotuner = Callable[[torch.nn.Module, int, torch.device], None] -_AUTOTUNERS: dict[str, Autotuner] = {} +_AUTOTUNERS: dict[str, tuple[Autotuner, int, frozenset[str] | None]] = {} -def register_autotuner(name: str, tuner: Autotuner) -> None: - """Register a kernel family's freeze-time launch-configuration tuner. +def register_autotuner( + name: str, + tuner: Autotuner, + min_level: int = 2, + lower_kinds: tuple[str, ...] | None = None, +) -> None: + """Register a kernel family's freeze-time tuner. Parameters ---------- name : str Unique family identifier; a repeated name replaces the prior tuner. tuner : Callable[[torch.nn.Module, int, torch.device], None] - Sweeps the shape keys ``model`` needs for this family that are not - already covered and registers the winners into the family's launch - table. Receives the resolved ``DP_TRITON_INFER`` level and the target - CUDA device. + Measures the keys ``model`` needs for this family that are not already + covered and registers the winners into the family's table. Receives + the resolved ``DP_TRITON_INFER`` level and the target CUDA device. + min_level : int, optional + Lowest ``DP_TRITON_INFER`` level at which the tuner runs, by default 2, + that of the Triton launch tables. + lower_kinds : tuple[str, ...] or None, optional + Lower kinds whose programs call the kernels of the family; a freeze + to another lower kind skips the tuner. By default it runs for every + lower kind. """ - _AUTOTUNERS[name] = tuner + _AUTOTUNERS[name] = ( + tuner, + min_level, + None if lower_kinds is None else frozenset(lower_kinds), + ) -def run_autotune(model: torch.nn.Module, device: torch.device) -> None: +def run_autotune(model: torch.nn.Module, device: torch.device, lower_kind: str) -> None: """Tune every registered kernel family for ``model`` on ``device``. Called by a freeze path after the model is built and before it is traced, - so ``resolve_*_config`` lookups made while tracing return freshly tuned - launches. A no-op below ``DP_TRITON_INFER`` level 2 or off CUDA. + so the lookups made while tracing return freshly tuned choices. A no-op + off CUDA; each tuner runs from the level it registered with, for the lower + kinds it registered with. Parameters ---------- @@ -69,19 +89,25 @@ def run_autotune(model: torch.nn.Module, device: torch.device) -> None: The model about to be frozen; tuners walk it for their shape keys. device : torch.device The target CUDA device the artifact will run on. + lower_kind : str + The lower kind of the program the freeze exports. """ level = triton_infer_level() - if level < 2 or device.type != "cuda" or not torch.cuda.is_available(): + if device.type != "cuda" or not torch.cuda.is_available(): return - for name, tuner in _AUTOTUNERS.items(): - # A tuning failure must degrade to the built-in / default launch tables - # rather than abort the freeze; the frozen artifact stays correct, only - # its launch configuration is untuned for uncovered shapes. + for name, (tuner, min_level, lower_kinds) in _AUTOTUNERS.items(): + if level < min_level or ( + lower_kinds is not None and lower_kind not in lower_kinds + ): + continue + # A tuning failure must degrade to the built-in / default tables rather + # than abort the freeze; the frozen artifact stays correct, only its + # launch choices are untuned for uncovered keys. try: tuner(model, level, device) except Exception: log.warning( - "Triton autotuner %r failed; keeping default launch configurations.", + "Autotuner %r failed; keeping the default launch choices.", name, exc_info=True, ) diff --git a/deepmd/pt_expt/kernels/cuda/dpa1/canonical.py b/deepmd/pt_expt/kernels/cuda/dpa1/canonical.py index b8b5c1a931..a6883787ba 100644 --- a/deepmd/pt_expt/kernels/cuda/dpa1/canonical.py +++ b/deepmd/pt_expt/kernels/cuda/dpa1/canonical.py @@ -182,6 +182,11 @@ def _generic_topology( source: torch.Tensor, destination_row_ptr: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_index_words, + ) + + source = canonical_index_words(source) physical_edge_count = int(destination_row_ptr[-1].item()) node_count = destination_row_ptr.shape[0] - 1 destination = torch.repeat_interleave( diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical.py b/deepmd/pt_expt/kernels/dpa4c/canonical.py index 750bb31b25..4c3f8eedd8 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical.py @@ -58,6 +58,21 @@ def canonical_model_eligible(model: Any) -> bool: return False if getattr(atomic_model, "atom_excl", None) is not None: return False + # A latent-charge long-range module joins the pipeline through the node + # positions and the cells, which its operator reads in place of a pair + # list or a mesh; that pipeline evaluates no pair potential. + if parts[0].has_long_range(): + from deepmd.pt_expt.long_range.latent_charge import ( + LatentChargeLongRange, + ) + + long_range = parts[0].long_range + if not ( + isinstance(long_range, LatentChargeLongRange) + and long_range.operator_eligible() + and parts[1] is None + ): + return False from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( mega_eligible, ) @@ -76,17 +91,35 @@ def op_available() -> bool: "dpa4c_canonical_compress", "dpa4c_canonical_compress_backward", "dpa4c_canonical_compress_backward_inplace", - "dpa4c_canonical_compress_energy_gradient", + "dpa4c_canonical_compress_energy_force", ) ) +def long_range_op_available() -> bool: + """Return whether the backend device carries the fused long-range operator.""" + return op_available() and operator_available( + "dpa4c_canonical_compress_long_range_energy_force" + ) + + +#: Position of the per-node magnetic moment among the compression arguments, +#: which open with the radial table (see ``compressed_operator_arguments``), +#: and among the inputs that follow ``edge_vec`` in the compact descriptor +#: operators, where the source indices, the destination row pointers and the +#: node types precede those arguments. +_SPIN_ARGUMENT_SLOT = 12 +_SPIN_SLOT = 3 + _SPIN_ARGUMENT_SLOT + + def _forward_fake( edge_vec: torch.Tensor, source: torch.Tensor, destination_row_ptr: torch.Tensor, atype: torch.Tensor, table: torch.Tensor, + basis_table: torch.Tensor, + basis_matrix: torch.Tensor, pair_film: torch.Tensor, pair_mixing: torch.Tensor, type_embedding: torch.Tensor, @@ -113,6 +146,8 @@ def _forward_fake( source, destination_row_ptr, table, + basis_table, + basis_matrix, pair_film, pair_mixing, readout_matrices, @@ -152,7 +187,7 @@ def _backward_fake( *args: Any, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: del descriptor_gradient, state - return _backward_shapes(edge_vec, args[13]) + return _backward_shapes(edge_vec, args[_SPIN_SLOT]) def _backward_inplace_fake( @@ -162,7 +197,7 @@ def _backward_inplace_fake( *args: Any, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: del descriptor_gradient, state - return _backward_shapes(edge_vec, args[13]) + return _backward_shapes(edge_vec, args[_SPIN_SLOT]) def _backward_shapes( @@ -187,40 +222,137 @@ def _backward_shapes( ) -def _energy_gradient_fake( +def _energy_force_fake( edge_vec: torch.Tensor, source: torch.Tensor, destination_row_ptr: torch.Tensor, + source_row_ptr: torch.Tensor, + source_order: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, atype: torch.Tensor, *args: Any, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - del source, destination_row_ptr - spin = args[10] +) -> tuple[torch.Tensor, ...]: + del source, destination_row_ptr, source_row_ptr, source_order, n_local + spin = args[_SPIN_ARGUMENT_SLOT] + nodes = atype.shape[0] + frames = n_node.shape[0] has_spin = spin.ndim == 2 - # Each absent output is allocated separately, because the schema declares - # four unannotated results and no two of them may share storage. + # Each output is allocated separately, because the schema declares seven + # unannotated results and no two of them may share storage. return ( - edge_vec.new_empty(atype.shape[0], 1, dtype=torch.float64), - torch.empty_like(edge_vec), - edge_vec.new_empty(atype.shape[0], 3, dtype=torch.float32) - if has_spin - else edge_vec.new_empty((0,), dtype=torch.float32), - torch.empty_like(edge_vec) + edge_vec.new_empty(frames, 1, dtype=torch.float64), + edge_vec.new_empty(nodes, 1, dtype=torch.float64), + edge_vec.new_empty(nodes, 3, dtype=torch.float64), + edge_vec.new_empty(frames, 3, 3, dtype=torch.float64), + edge_vec.new_empty(nodes, 3, 3, dtype=torch.float32), + edge_vec.new_empty(nodes, 3, dtype=torch.float64) if has_spin else edge_vec.new_empty((0,), dtype=torch.float32), + edge_vec.new_empty(nodes, dtype=torch.int32), + ) + + +def _long_range_energy_force_fake( + edge_vec: torch.Tensor, + source: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_row_ptr: torch.Tensor, + source_order: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + atype: torch.Tensor, + *args: Any, +) -> tuple[torch.Tensor, ...]: + del source, destination_row_ptr, source_row_ptr, source_order, n_local, args + nodes = atype.shape[0] + frames = n_node.shape[0] + # Each output is allocated separately, because the schema declares six + # unannotated results and no two of them may share storage. + return ( + edge_vec.new_empty(frames, 1, dtype=torch.float64), + edge_vec.new_empty(nodes, 1, dtype=torch.float64), + edge_vec.new_empty(nodes, 3, dtype=torch.float64), + edge_vec.new_empty(frames, 3, 3, dtype=torch.float64), + edge_vec.new_empty(nodes, 3, 3, dtype=torch.float32), + edge_vec.new_empty(nodes, dtype=torch.int32), ) -def _cpu_energy_gradient( +def _cpu_energy_force( + edge_vec: torch.Tensor, + source: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_row_ptr: torch.Tensor, + source_order: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + atype: torch.Tensor, *args: Any, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Reference sequence of the fused operator, evaluated in one run.""" - descriptor_args = args[:_DESCRIPTOR_ARGUMENT_COUNT] - ws, bs, resnets, w_head, b_head, bias_atom_e, act, seed, _tile, pair_table = args[ - _DESCRIPTOR_ARGUMENT_COUNT: +) -> tuple[torch.Tensor, ...]: + """Reference sequence of the fused operator, evaluated in one run. + + The force and the virial are assembled by index reduction in fp64. The + per-atom virial is attributed to the source nodes that the source CSR + lists, and is zero when two-entry placeholders take the place of the CSR. + A frame whose energy or virial is not finite reports every force of its + nodes, and every magnetic force, as not a number, as the CUDA operator + does. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + frame_id_from_n_node, + node_ownership_mask, + ) + + node_count = atype.shape[0] + edge_storage = edge_vec.shape[0] + device = edge_vec.device + complete = ( + source_row_ptr.numel() == node_count + 1 + and source_order.numel() == edge_storage + ) + if not complete and not (source_row_ptr.numel() == 2 and source_order.numel() == 2): + raise RuntimeError( + "dpa4c_canonical_compress_energy_force: the source CSR must hold " + f"{node_count + 1} row pointers and {edge_storage} positions, or two " + "entries each to decline the per-atom virial; got " + f"{source_row_ptr.numel()} and {source_order.numel()}" + ) + if node_count == 0: + # An empty node axis holds no energy, force or virial: the outputs + # take the layout of the operator with zero frame totals, as in the + # CUDA operator. + outputs = _energy_force_fake( + edge_vec, + source, + destination_row_ptr, + source_row_ptr, + source_order, + n_node, + n_local, + atype, + *args, + ) + outputs[0].zero_() + outputs[3].zero_() + return outputs + real = atype >= 0 + mask = real & node_ownership_mask(n_node, n_local, node_count) + atype = torch.where(real, atype, torch.zeros_like(atype)) + seed = mask.to(torch.float64) + descriptor_args = ( + edge_vec, + source, + destination_row_ptr, + atype, + *args[:_DESCRIPTOR_TAIL_COUNT], + ) + # The trailing fitting-route table, which a caller may leave to its default, + # sets only the speed of the CUDA operator. + ws, bs, resnets, w_head, b_head, bias_atom_e, act, _tile, pair_table = args[ + _DESCRIPTOR_TAIL_COUNT : _DESCRIPTOR_TAIL_COUNT + 9 ] descriptor, state = _cpu_forward(*descriptor_args) - atype = descriptor_args[3] energy, saved = torch.ops.deepmd.graph_fitting( descriptor, atype, ws, bs, resnets, w_head, b_head, bias_atom_e, act ) @@ -232,14 +364,86 @@ def _cpu_energy_gradient( ) if pair_table.numel(): energy = energy + pair_energy - return energy, edge_gradient, spin_gradient, edge_spin_gradient + atom_energy = energy * seed[:, None] + + edge_index, edge_mask, _ = _generic_topology(source, destination_row_ptr) + keep = edge_mask[:, None].to(torch.float64) + g = edge_gradient.to(torch.float64) * keep + outer = -(g[:, :, None] * edge_vec.to(torch.float64)[:, None, :]) + source_index, destination_index = edge_index[0], edge_index[1] + destination_force = torch.zeros( + node_count, 3, dtype=torch.float64, device=device + ).index_add_(0, destination_index, g) + force = destination_force.clone().index_add_(0, source_index, -g) + # A force contribution that lies outside the range of the fixed-point totals + # of the CUDA operator -- an edge onto its source, the row of a destination, + # or the magnetic contribution of an edge onto its source -- makes its node + # report its energy as not a number, as the CUDA operator does. + out_of_range = (destination_force.abs() >= _FORCE_FIXED_LIMIT).any(1) + out_of_range[source_index[(g.abs() >= _FORCE_FIXED_LIMIT).any(1)]] = True + if spin_gradient.ndim == 2: + spin_edges = (edge_spin_gradient.abs() * keep >= _FORCE_FIXED_LIMIT).any(1) + out_of_range[source_index[spin_edges]] = True + atom_energy[out_of_range] = torch.nan + node_virial = torch.zeros( + node_count, 3, 3, dtype=torch.float64, device=device + ).index_add_(0, destination_index, outer) + # Padding nodes past the last frame belong to none: they enter no frame + # total, as in the CUDA operator. + frame_id = frame_id_from_n_node(n_node, n_total=node_count) + in_frame = torch.arange(node_count, dtype=torch.int64, device=device) < n_node.sum() + virial = torch.zeros( + n_node.shape[0], 3, 3, dtype=torch.float64, device=device + ).index_add_(0, frame_id[in_frame], node_virial[in_frame]) + frame_energy = torch.zeros( + n_node.shape[0], 1, dtype=torch.float64, device=device + ).index_add_(0, frame_id[in_frame], atom_energy[in_frame]) + atom_virial = torch.zeros(node_count, 3, 3, dtype=torch.float64, device=device) + if complete: + listed = source_row_ptr[1:] - source_row_ptr[:-1] + rows = torch.repeat_interleave( + torch.arange(node_count, dtype=torch.int64, device=device), listed + ) + positions = source_order[: int(source_row_ptr[-1])].to(torch.int64) + atom_virial.index_add_(0, rows, outer[positions]) + if spin_gradient.ndim == 2: + spin_force = torch.zeros( + node_count, 3, dtype=torch.float64, device=device + ).index_add_(0, source_index, edge_spin_gradient.to(torch.float64) * keep) + force_mag = -(spin_gradient.to(torch.float64) + spin_force) + else: + force_mag = spin_gradient.new_empty((0,), dtype=torch.float32) + # A fixed-point force cannot carry a non-finite value through the CUDA + # operator, which marks the forces of a broken frame instead; padding + # nodes past the last frame belong to none. + broken = ~( + torch.isfinite(frame_energy).all(1) & torch.isfinite(virial).flatten(1).all(1) + ) + node_broken = broken[frame_id] & in_frame + force[node_broken] = torch.nan + if force_mag.ndim == 2: + force_mag[node_broken] = torch.nan + return ( + frame_energy, + atom_energy, + force, + virial, + atom_virial.to(torch.float32), + force_mag, + mask.to(torch.int32), + ) def _generic_topology( source: torch.Tensor, destination_row_ptr: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Materialize the generic topology for the small CPU trace sample.""" + """Materialize the generic topology of a compact graph for the CPU paths.""" + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_index_words, + ) + + source = canonical_index_words(source) physical_edge_count = int(destination_row_ptr[-1].item()) node_count = destination_row_ptr.shape[0] - 1 destination = torch.repeat_interleave( @@ -267,10 +471,19 @@ def _generic_topology( # leading ``canonical`` flag; the remaining trailing scalars are identical. _CANONICAL_SCALAR_COUNT = 8 -#: Leading arguments of the fused operator that describe the descriptor: -#: ``edge_vec`` plus the compact topology, the compression artifacts and the -#: eight trailing geometry scalars. -_DESCRIPTOR_ARGUMENT_COUNT = 26 +#: Arguments of the compact descriptor operator: ``edge_vec``, the source +#: indices, the destination row pointers and the node types, then the +#: compression artifacts and the trailing geometry scalars. +_DESCRIPTOR_ARGUMENT_COUNT = 28 + +#: Magnitude in eV/Å from which a force contribution lies outside the range of +#: the fixed-point force totals of the CUDA operator (``kForceFixedLimit``). +_FORCE_FIXED_LIMIT = 2.0**30 + +#: Descriptor arguments that follow the node types. The fused energy-force +#: operator passes them after its own graph inputs, which insert the source +#: CSR and the frame counts before the node types. +_DESCRIPTOR_TAIL_COUNT = _DESCRIPTOR_ARGUMENT_COUNT - 4 def _cpu_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: @@ -354,9 +567,13 @@ def ensure_registered() -> None: torch.library.register_fake("deepmd::dpa4c_canonical_compress_backward")( _backward_fake ) - torch.library.register_fake("deepmd::dpa4c_canonical_compress_energy_gradient")( - _energy_gradient_fake + torch.library.register_fake("deepmd::dpa4c_canonical_compress_energy_force")( + _energy_force_fake ) + if long_range_op_available(): + torch.library.register_fake( + "deepmd::dpa4c_canonical_compress_long_range_energy_force" + )(_long_range_energy_force_fake) torch.library.register_fake("deepmd::dpa4c_canonical_compress_backward_inplace")( _backward_inplace_fake ) @@ -368,8 +585,8 @@ def ensure_registered() -> None: "CPU", ) _cpu_library.impl( - "dpa4c_canonical_compress_energy_gradient", - _cpu_energy_gradient, + "dpa4c_canonical_compress_energy_force", + _cpu_energy_force, "CPU", ) _cpu_library.impl( @@ -384,11 +601,10 @@ def dpa4c_canonical_compress_energy_force( fitting: Any, graph: CanonicalGraph, atype: torch.Tensor, - ownership: torch.Tensor, atom_bias: torch.Tensor, - do_atomic_virial: bool, spin: torch.Tensor | None = None, pair_table: torch.Tensor | None = None, + fitting_routes: list[int] | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -396,13 +612,28 @@ def dpa4c_canonical_compress_energy_force( torch.Tensor, torch.Tensor, torch.Tensor, + torch.Tensor, ]: """Evaluate compressed DPA4C from a compact canonical edge stream. The compact ABI carries only source indices and CSR row pointers, so the descriptor operator addresses edges by position and omits every mask, padding-type, and table-tail check that the cutoff-compacted graph already - guarantees. The backward reuses the forward state in place. + guarantees. The backward reuses the forward state in place and assembles + the force and the virial itself, so the source CSR is read only for the + per-atom virial: a graph that carries placeholders in its place yields a + zero per-atom virial. The operator also derives the ownership mask, masks + the node energies and forms the frame totals, so that nothing around it + runs on the node axis. + + The forces are summed in fixed point. A frame whose energy or virial is + not finite reports every force of its nodes, and every magnetic force, as + not a number, and so does a frame with a force contribution of 2^30 eV/Å + or more, which lies outside the range of the sums; its node then reports + its energy as not a number. The check covers each contribution, so a sum + of three or more smaller ones beyond the range of 2^31 eV/Å, which no + configuration short of three atoms within about 0.01 Å of one atom + reaches, is not detected. Parameters ---------- @@ -411,16 +642,16 @@ def dpa4c_canonical_compress_energy_force( fitting Eligible pt_expt energy fitting network. graph - Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` and - whose ``destination_row_ptr`` has shape ``(N + 1,)`` in ``int64``. + Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` or + ``int64``, or shape ``(E, W)`` in ``uint8`` for the exported bytes of + those words, and whose ``destination_row_ptr`` has shape ``(N + 1,)`` + in ``int64``. Its source CSR is either complete or a two-entry + placeholder. atype - Flat node atom types with shape ``(N,)``. - ownership - Boolean mask selecting energy-contributing nodes with shape ``(N,)``. + Flat node atom types with shape ``(N,)``; a negative type marks a + virtual node, which contributes no energy. atom_bias Combined atomic energy bias with shape ``(ntypes,)`` in eV. - do_atomic_virial - Whether to return per-node virials. spin Per-node magnetic moments with shape ``(N, 3)`` for a spin-conditioned descriptor, or ``None``. @@ -428,44 +659,44 @@ def dpa4c_canonical_compress_energy_force( Series constants of an analytical pair potential with shape ``((T + 1) ** 2, 8)``, or ``None``. The fused operator adds its energy to the node energies and its slope to the edge gradient. + fitting_routes + Fitting-route table of the CUDA operator, which sets only its speed + (``deepmd.pt_expt.kernels.fitting_routes``). ``None`` or an empty + table leaves the routes to the rules of the operator. Returns ------- energy Per-frame energy with shape ``(F, 1)`` in eV, fp64. atom_energy - Per-node energy with shape ``(N, 1)`` in eV, fp64. + Per-node energy with shape ``(N, 1)`` in eV, fp64, zero for every node + outside the ownership mask whose forces lie in range. force - Per-node force with shape ``(N, 3)`` in eV/Å, fp32. + Per-node force with shape ``(N, 3)`` in eV/Å, fp64. virial - Per-frame virial with shape ``(F, 3, 3)`` in eV, fp32. + Per-frame virial with shape ``(F, 3, 3)`` in eV, fp64. atom_virial - Per-node virial with shape ``(N, 3, 3)`` in eV, or an empty tensor. + Per-node virial with shape ``(N, 3, 3)`` in eV, fp32, attributed to + source nodes; zero when the graph carries no source CSR. force_mag - Per-node magnetic force with shape ``(N, 3)``, or an empty tensor. The - on-site part closes inside the fused operator; the neighbour part is - emitted per edge and reduced onto source nodes here, where the source - CSR is in scope. + Per-node magnetic force with shape ``(N, 3)`` in fp64 for a + spin-conditioned descriptor, or an empty tensor. Both the on-site part + and the neighbour part close inside the fused operator. + mask + Ownership mask with shape ``(N,)``, int32: one for a real node among + the leading ``n_local`` nodes of its frame. Raises ------ ValueError - If the model or the compiled operators do not support the compact path. + If the descriptor lies outside the configurations the compiled + operator specializes on, or the fused fitting operator cannot + reproduce the fitting network. """ from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( compressed_operator_arguments, mega_eligible, ) - from deepmd.pt_expt.kernels.edge_force_virial import ( - canonical_edge_force_virial, - canonical_op_available, - ) - from deepmd.pt_expt.kernels.edge_force_virial import ( - ensure_registered as ensure_force_registered, - ) - from deepmd.pt_expt.kernels.edge_force_virial import ( - frame_scalar_sum, - ) from deepmd.pt_expt.kernels.graph_fitting import ( ensure_registered as ensure_fitting_registered, ) @@ -476,20 +707,18 @@ def dpa4c_canonical_compress_energy_force( ensure_registered() ensure_fitting_registered() - ensure_force_registered() - if not mega_eligible(descriptor) or not canonical_op_available(): + if not mega_eligible(descriptor): raise ValueError("model is not eligible for compact canonical DPA4C inference") network = fitting_operator_arguments(fitting) - ( - atom_energy_raw, - edge_gradient, - spin_gradient, - edge_spin_gradient, - ) = torch.ops.deepmd.dpa4c_canonical_compress_energy_gradient( + return torch.ops.deepmd.dpa4c_canonical_compress_energy_force( graph.edge_vec, graph.source, graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + graph.n_node, + graph.n_local, atype, *compressed_operator_arguments(descriptor, spin), int(descriptor.lmax), @@ -501,24 +730,10 @@ def dpa4c_canonical_compress_energy_force( network.head_bias, atom_bias.to(torch.float64).contiguous(), network.activation, - ownership.to(torch.float64).reshape(-1).contiguous(), node_tile(), graph.edge_vec.new_empty(0) if pair_table is None else pair_table, + # No table is passed as None rather than as an empty list: the + # AOTInductor proxy executor of torch 2.13 and earlier hands a + # serialized empty int list to the operator as None. + list(fitting_routes) if fitting_routes else None, ) - atom_energy = atom_energy_raw * ownership[:, None].to(atom_energy_raw.dtype) - energy = frame_scalar_sum(atom_energy, graph.n_node) - # The magnetic cotangent is reduced onto its source nodes by the force - # assembly, which already walks that grouping. - force, atom_virial, virial, source_spin = canonical_edge_force_virial( - edge_gradient, - graph.edge_vec, - graph.destination_row_ptr, - graph.source_row_ptr, - graph.source_order, - graph.n_node, - edge_spin_gradient, - atype.shape[0], - do_atomic_virial, - ) - force_mag = spin_gradient if spin is None else -(spin_gradient + source_spin) - return energy, atom_energy, force, virial, atom_virial, force_mag diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py new file mode 100644 index 0000000000..b588eb04ef --- /dev/null +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compact canonical deployment of compressed DPA4C with a long-range module.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch + +if TYPE_CHECKING: + from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + ) + + +def fused_long_range_eligible(long_range: Any) -> bool: + """Whether the fused long-range operator evaluates a long-range module. + + The operator evaluates a charge network of one channel that the fused + fitting kernels reproduce, beside the compressed descriptor and the energy + fitting that the compact canonical lower already requires; a module whose + charge network they do not reproduce, or a library without the operator, + takes :func:`dpa4c_canonical_long_range_energy_force`. + """ + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + long_range_op_available, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_eligible, + ) + + return long_range_op_available() and fitting_eligible(long_range.charge_net) + + +def dpa4c_canonical_compress_long_range_energy_force( + descriptor: Any, + fitting: Any, + long_range: Any, + graph: CanonicalGraph, + atype: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor | None, + box: torch.Tensor | None, + atom_bias: torch.Tensor, + fitting_routes: list[int] | None = None, + charge_fitting_routes: list[int] | None = None, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Evaluate compressed DPA4C and its long-range module in one operator. + + The operator evaluates the descriptor of every node, the charges of the + module from it, the energy fitting, the long range of the charges and one + descriptor backward from the cotangents of both networks, and assembles + the forces in fixed point like + :func:`~deepmd.pt_expt.kernels.dpa4c.canonical.dpa4c_canonical_compress_energy_force`, + so the force assembly reads the source CSR only for the per-atom virial. + It derives the ownership mask from the frame counts, and the charged + nodes of the module are the owned real nodes. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor without a spin condition. + fitting + Eligible pt_expt energy fitting network. + long_range + Long-range module for which :func:`fused_long_range_eligible` holds. + graph + Canonical graph as for + :func:`~deepmd.pt_expt.kernels.dpa4c.canonical.dpa4c_canonical_compress_energy_force`; + its source CSR is either complete or a two-entry placeholder, which + declines the per-atom virial. + atype + Flat node atom types with shape ``(N,)``; a negative type marks a + virtual node, which contributes no energy and carries no charge. + coord + Node positions with shape ``(N, 3)`` in Å. + communicator + The address of the MPI_Comm handle of a frame decomposed into + domains, as an int64 CPU tensor of one element, or zero or ``None`` + for a single domain. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the CPU, a + zero row marking a complete non-periodic frame, or ``None`` when no + frame is periodic. + atom_bias + Combined atomic energy bias with shape ``(ntypes,)`` in eV. + fitting_routes + Fitting-route table of the energy fitting, as for the fused operator + without a long-range module. + charge_fitting_routes + Fitting-route table of the charge network. ``None`` or an empty table + leaves the routes of its network to the rules of the operator, for + either table. + + Returns + ------- + energy + Per-frame energy with shape ``(F, 1)`` in eV, fp64. + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, fp64. + force + Per-node force with shape ``(N, 3)`` in eV/Å, fp64. + virial + Per-frame virial with shape ``(F, 3, 3)`` in eV, fp64. + atom_virial + Per-node virial with shape ``(N, 3, 3)`` in eV, fp32; zero when the + graph carries no source CSR. + mask + Ownership mask with shape ``(N,)``, int32. + """ + from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, + ) + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + ensure_registered, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + compressed_operator_arguments, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + ensure_registered as ensure_fitting_registered, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_operator_arguments, + node_tile, + ) + from deepmd.pt_expt.kernels.long_range import ( + kernel_operator_arguments, + ) + + ensure_registered() + ensure_fitting_registered() + network = fitting_operator_arguments(fitting) + charge_net = long_range.charge_net + charges = fitting_operator_arguments(charge_net) + total_charge, channel = long_range.operator_constraint(graph.n_node.shape[0], coord) + kernel, coefficient = kernel_operator_arguments(long_range.kernel) + if communicator is None: + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") + return torch.ops.deepmd.dpa4c_canonical_compress_long_range_energy_force( + graph.edge_vec, + graph.source, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + graph.n_node, + graph.n_local, + atype, + *compressed_operator_arguments(descriptor), + int(descriptor.lmax), + *descriptor._compression_scalars, + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.head_bias, + atom_bias.to(torch.float64).contiguous(), + network.activation, + node_tile(), + charges.weights, + charges.biases, + charges.residuals, + charges.head_weight, + charges.head_bias, + charge_net.bias_atom_e[:, 0].to(torch.float64).contiguous(), + charges.activation, + coord.to(torch.float64).contiguous(), + total_charge, + # A module that does not serve periodic systems sums every frame + # directly. + box if long_range.periodic else None, + coefficient.to(coord.device), + communicator, + kernel, + channel, + COULOMB_CONSTANT, + long_range.pme_order, + long_range.mesh_spacing() if long_range.periodic else 1.0, + # No table is passed as None rather than as an empty list: the + # AOTInductor proxy executor of torch 2.13 and earlier hands a + # serialized empty int list to the operator as None. + list(fitting_routes) if fitting_routes else None, + list(charge_fitting_routes) if charge_fitting_routes else None, + ) + + +def dpa4c_canonical_long_range_energy_force( + descriptor: Any, + fitting: Any, + long_range: Any, + graph: CanonicalGraph, + atype: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor | None, + box: torch.Tensor | None, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate compressed DPA4C and its long-range module from a compact graph. + + The descriptor operator runs once, and its output feeds both the fused + energy fitting and the charge network of the long-range module. The + derivative of the energy with respect to the descriptor collects the + fitting term and the response of the charges, so a single descriptor + backward yields the edge gradient of both; the long-range force at fixed + charges is added to the assembled force. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor without a spin condition. + fitting + Eligible pt_expt energy fitting network. + long_range + Long-range module whose ``operator_eligible()`` holds. + graph + Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` or + ``int64``, or shape ``(E, W)`` in ``uint8`` for the exported bytes of + those words, and whose ``destination_row_ptr`` has shape ``(N + 1,)`` + in ``int64``. Its source CSR is complete. + atype + Flat node atom types with shape ``(N,)``. + coord + Node positions with shape ``(N, 3)`` in Å. + communicator + The address of the MPI_Comm handle of a frame decomposed into + domains, one per rank, as an int64 CPU tensor of one element, or zero + or ``None`` for a single domain. Across domains the owned nodes of + every rank make up the frame. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the CPU, + a zero row marking a complete non-periodic frame, or ``None`` when no + frame is periodic. + ownership + Boolean mask selecting energy-contributing nodes with shape ``(N,)``. + atom_bias + Combined atomic energy bias with shape ``(ntypes,)`` in eV. + do_atomic_virial + Whether to return per-node virials. + + Returns + ------- + energy + Per-frame energy with shape ``(F, 1)`` in eV, fp64. + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, fp64. + force + Per-node force with shape ``(N, 3)`` in eV/Å, fp32. + virial + Per-frame virial with shape ``(F, 3, 3)`` in eV, fp32. + atom_virial + Per-node virial with shape ``(N, 3, 3)`` in eV, or an empty tensor. + """ + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + ensure_registered, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + compressed_operator_arguments, + fitting_energy_and_gradient, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + canonical_edge_force_virial, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + ensure_registered as ensure_force_registered, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + frame_scalar_sum, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + ensure_registered as ensure_fitting_registered, + ) + + ensure_registered() + ensure_fitting_registered() + ensure_force_registered() + operator_args = ( + graph.source, + graph.destination_row_ptr, + atype, + *compressed_operator_arguments(descriptor), + int(descriptor.lmax), + *descriptor._compression_scalars, + ) + node_descriptor, state = torch.ops.deepmd.dpa4c_canonical_compress( + graph.edge_vec, *operator_args + ) + # The charge network reads the descriptor before the fused fitting + # overwrites it with its own cotangent. + ( + long_range_energy, + long_range_gradient, + long_range_force, + long_range_atom_virial, + long_range_virial, + ) = long_range.energy_force_from_coord( + node_descriptor, atype, ownership, graph.n_node, coord, communicator, box + ) + atom_energy, descriptor_gradient = fitting_energy_and_gradient( + fitting, node_descriptor, atype, ownership, atom_bias + ) + edge_gradient, _, _ = torch.ops.deepmd.dpa4c_canonical_compress_backward_inplace( + descriptor_gradient + long_range_gradient, + state, + graph.edge_vec, + *operator_args, + ) + force, atom_virial, virial, _ = canonical_edge_force_virial( + edge_gradient, + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + graph.n_node, + graph.edge_vec.new_empty(0), + atype.shape[0], + do_atomic_virial, + ) + # Nodes outside the ownership carry no charge and hence no long-range + # energy; the halo nodes of a periodic frame are images, which the + # reciprocal-space sum holds itself. + atom_energy = atom_energy + long_range_energy + energy = frame_scalar_sum(atom_energy, graph.n_node) + virial = virial + long_range_virial.reshape(-1, 3, 3) + if do_atomic_virial: + atom_virial = atom_virial + long_range_atom_virial.reshape(-1, 3, 3) + return energy, atom_energy, force + long_range_force, virial, atom_virial diff --git a/deepmd/pt_expt/kernels/dpa4c/graph_compress.py b/deepmd/pt_expt/kernels/dpa4c/graph_compress.py index a7805d5c1a..8e9532280b 100644 --- a/deepmd/pt_expt/kernels/dpa4c/graph_compress.py +++ b/deepmd/pt_expt/kernels/dpa4c/graph_compress.py @@ -21,6 +21,13 @@ mode-mixing caches, the readout projections, and the sparse angular coupling tables are materialized once when compression is enabled. +The scalar-only channels from :math:`C_1` on feed degree zero alone, so a +trained radial map, which spans far fewer functions than it has such +channels, compresses further: :func:`build_low_rank_table` tabulates a +low-rank basis of them, which the single-type scans and the scans of a system +of a few types read in place of the channels, when the basis reproduces every +channel at the knots. + Degrees one and two dominate the readout and are contracted in closed form. Degrees three and four carry a single channel each, so their couplings are driven by a compact sparse Cartesian Gaunt table rather than by specialized @@ -91,7 +98,9 @@ "ChargeStateFold", "build_charge_state_artifacts", "build_compression_artifacts", + "build_low_rank_table", "build_radial_table", + "compression_ineligibility_reason", "coupling_records", "descriptor_profile", "dpa4c_graph_compress", @@ -128,6 +137,25 @@ # below the retained coupling entries, far below the fp32 working precision. _COUPLING_TOLERANCE = 1.0e-9 +# A low-rank radial basis is adopted when it reproduces every scalar-only +# channel and its radial derivative at the table knots to this fraction of +# their largest magnitudes, below the single-precision rounding of the +# tabulated coefficients. +_LOW_RANK_TOLERANCE = 1.0e-8 + +# Length in Å over which the radial derivatives enter the fit of that basis, +# so that a channel value and its change over this length weigh alike. +_LOW_RANK_SLOPE_LENGTH = 1.0 + +# Width of the low-rank expansion per channel count, mirroring +# ``low_rank_basis`` on the device side; absent widths admit none. +_LOW_RANK_BASIS = {64: 32, 128: 32} + +# Mode profiles a folded single-type scan carries, mirroring ``kFoldedModes`` +# on the device side. A descriptor with more modes scans unfolded and never +# reads a low-rank table. +_FOLDED_MODES = 4 + # Unit-Frobenius Cartesian normalizations of the two closed-form couplings. _INV_SQRT_FIVE = 1.0 / math.sqrt(5.0) _BIS222_SCALE = math.sqrt(12.0 / 35.0) @@ -212,6 +240,55 @@ def mega_eligible(descriptor: Any) -> bool: ) +def compression_ineligibility_reason(descriptor: Any) -> str | None: + """Return why a DPA4C descriptor cannot build a compressed snapshot. + + The same device-independent conditions govern explicit compression and + automatic deployment. Kernel availability is checked by the inference + dispatcher rather than by artifact construction. + + Parameters + ---------- + descriptor + DPA4C descriptor with its configuration and parameters loaded. + + Returns + ------- + str or None + Diagnostic for an unsupported configuration or initial charge state, + or ``None`` when these compression prerequisites are satisfied. + """ + if descriptor.exclude_types: + return ( + "DPA4C compressed inference has no type-exclusion branch, so a " + "descriptor with a non-empty `exclude_types` cannot reach the " + "fused path. Drop the exclusions, or deploy the uncompressed " + "graph archive." + ) + if descriptor.stddev.dtype != torch.float32: + return ( + "DPA4C compression requires descriptor precision `float32`, " + f"got {descriptor.stddev.dtype}" + ) + if not mega_eligible(descriptor): + return ( + "DPA4C compression supports " + f"channels {SUPPORTED_CHANNELS}, lmax {SUPPORTED_LMAX} and " + f"radial_modes {SUPPORTED_RADIAL_MODES}, got " + f"channels={descriptor.channels}, lmax={descriptor.lmax}, " + f"radial_modes={descriptor.radial_modes}" + ) + if ( + descriptor.charge_spin_embedding is not None + and descriptor.default_chg_spin is None + ): + return ( + "A charge-conditioned DPA4C snapshot requires `default_chg_spin` " + "to supply its initial frame condition." + ) + return None + + # === Compiled profile === @@ -239,6 +316,9 @@ class DescriptorProfile: Descriptor coordinate of the first bispectrum entry. spin_width Width of the trailing native-spin invariant block. + basis + Width of the low-rank expansion of the scalar-only channels, zero for + a profile that admits no expansion. """ channels: int @@ -251,6 +331,7 @@ class DescriptorProfile: bispectrum_base: int spin_channels: int spin_width: int + basis: int @property def state_width(self) -> int: @@ -348,6 +429,8 @@ def descriptor_profile( bispectrum_base=bispectrum_base, spin_channels=spin_channels, spin_width=spin_dim, + # The spin branch never scans folded, so it carries no expansion. + basis=0 if has_spin else _LOW_RANK_BASIS.get(int(channels), 0), ) @@ -546,30 +629,29 @@ def _quintic_coefficients( ) -def build_radial_table( +def _radial_knots( descriptor: Any, - stride: float = 0.002, -) -> tuple[torch.Tensor, torch.Tensor]: - """Tabulate the composed DPA4C radial branch and its mode profiles. - - The table width is ``channels + radial_modes``: the leading block holds the - shared radial map read by every degree, and the trailing block holds the - shared mode profiles that each ordered type pair mixes. + stride: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate every tabulated distance map at the knots of the radial table. Parameters ---------- descriptor - pt_expt DPA4C descriptor whose DPA4 radial modules define the table. + pt_expt DPA4C descriptor whose DPA4 radial modules define the maps. stride Uniform distance spacing in Å. Returns ------- - table - Quintic coefficients with shape ``(S, 6 * (channels + radial_modes))``, - fp32. - info - CPU metadata ``[stride, table_max, rcut, eps, degree_floor]``, fp64. + values + Channel maps followed by the mode profiles at the knots + ``0, stride, ..., S * stride``, with shape + ``(S + 1, channels + radial_modes)``, fp64. + first + Their first radial derivatives in Å⁻¹, with the same shape, fp64. + second + Their second radial derivatives in Å⁻², with the same shape, fp64. Raises ------ @@ -596,7 +678,6 @@ def build_radial_table( ) ) interval_count = math.ceil(float(descriptor.rcut) / float(stride)) - table_max = interval_count * float(stride) distance = torch.arange( interval_count + 1, dtype=torch.float64, @@ -627,16 +708,58 @@ def build_radial_table( ) first_columns.append(first_channel) second_columns.append(second_channel) - first = torch.stack(first_columns, dim=-1) - second = torch.stack(second_columns, dim=-1) + return ( + values.detach(), + torch.stack(first_columns, dim=-1).detach(), + torch.stack(second_columns, dim=-1).detach(), + ) + - # === Step 3. Convert knot data to the runtime spline layout === +def build_radial_table( + descriptor: Any, + stride: float = 0.002, + *, + knots: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Tabulate the composed DPA4C radial branch and its mode profiles. + + The table width is ``channels + radial_modes``: the leading block holds the + shared radial map read by every degree, and the trailing block holds the + shared mode profiles that each ordered type pair mixes. + + Parameters + ---------- + descriptor + pt_expt DPA4C descriptor whose DPA4 radial modules define the table. + stride + Uniform distance spacing in Å. + knots + The knot data of the maps at this stride, as ``_radial_knots`` returns + it, or ``None`` to evaluate it here. + + Returns + ------- + table + Quintic coefficients with shape ``(S, 6 * (channels + radial_modes))``, + fp32. + info + CPU metadata ``[stride, table_max, rcut, eps, degree_floor]``, fp64. + + Raises + ------ + ValueError + If ``stride`` is not positive. + """ + if stride <= 0.0: + raise ValueError(f"`stride` must be positive, got {stride}") + values, first, second = ( + _radial_knots(descriptor, stride) if knots is None else knots + ) table = _quintic_coefficients(values, first, second, float(stride)) - table = table.detach().to(torch.float32).contiguous() info = torch.tensor( [ float(stride), - table_max, + (values.shape[0] - 1) * float(stride), float(descriptor.rcut), float(descriptor._EPS), float(descriptor._DEGREE_NORM_FLOOR), @@ -644,7 +767,153 @@ def build_radial_table( dtype=torch.float64, device="cpu", ) - return table, info + return table.to(torch.float32).contiguous(), info + + +def build_low_rank_table( + descriptor: Any, + stride: float = 0.002, + *, + knots: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r"""Tabulate a low-rank basis of the scalar-only radial channels. + + The channels from :math:`C_1` on feed degree zero alone, so their moments + are linear in the channel values: a single-type scan may accumulate the + moments of :math:`K` basis functions instead and expand them into the + channel moments once per node, and the scan of a system of a few types + once per node and type of the system. Each such channel is divided by its + largest magnitude over the knots, and the basis is the leading right + singular subspace of the matrix that stacks the scaled values and the + scaled radial derivatives over ``_LOW_RANK_SLOPE_LENGTH`` at every knot; + by the Eckart-Young theorem it is the best rank-:math:`K` approximation of + the channels in that norm. A basis function is a fixed linear combination + of channels, so its knot values and derivatives, and therefore its quintic + coefficients, follow exactly from those of the channels. + + The basis is adopted only when it reproduces, in fp64 and at every knot, + every scalar-only channel and its radial derivative to + ``_LOW_RANK_TOLERANCE`` of their largest magnitudes. That tolerance lies + below the rounding of the single-precision coefficients, so the basis adds + no error of its own to the knot data the check covers; the second + derivatives are projected onto the same basis but not checked, and within + an interval they enter the interpolant at the order of the squared stride. + A table whose knots give no more samples than the basis has functions + lets any channels pass the check, since the basis then spans every sample, + so such a table keeps the direct one, and so does a table whose knot data + is not finite. When no basis is adopted, for a profile that admits no + expansion, and for a descriptor with more mode profiles than the folded + scans carry, both returned tensors are empty and the scan reads the + channels directly. + + The basis functions are stored in increasing order of their singular + values. The expansion sums them in storage order, so the smallest + contributions enter first and the rounding of its partial sums stays at + the scale of the channel they converge to. + + Parameters + ---------- + descriptor + pt_expt DPA4C descriptor whose DPA4 radial modules define the table. + stride + Uniform distance spacing in Å, equal to that of the radial table. + knots + The knot data of the maps at this stride, as ``_radial_knots`` returns + it, or ``None`` to evaluate it here. + + Returns + ------- + basis_table + Quintic coefficients with shape ``(S, 6 * (C1 + K + radial_modes))``, + fp32: the degree-one channels, the basis functions and the mode + profiles, in the layout and on the knots of the radial table. Empty + when no basis is adopted. + basis_matrix + Expansion with shape ``(K, channels - C1)``, fp32: channel ``C1 + c`` + is the sum over ``k`` of ``basis_matrix[k, c]`` times basis function + ``k``. Empty when no basis is adopted. + + Raises + ------ + ValueError + If ``stride`` is not positive. + """ + if stride <= 0.0: + raise ValueError(f"`stride` must be positive, got {stride}") + profile = descriptor_profile( + descriptor.channels, + descriptor.lmax, + descriptor.spin is not None, + ) + + def empty() -> tuple[torch.Tensor, torch.Tensor]: + # Two distinct tensors, which the descriptor stores as two buffers. + device = descriptor.stddev.device + return ( + torch.empty(0, dtype=torch.float32, device=device), + torch.empty(0, dtype=torch.float32, device=device), + ) + + if profile.basis == 0 or int(descriptor.radial_modes) > _FOLDED_MODES: + return empty() + values, first, second = ( + _radial_knots(descriptor, stride) if knots is None else knots + ) + split = profile.degree_channels[1] + scalar = slice(split, profile.channels) + tiny = torch.finfo(torch.float64).tiny + value_scale = values[:, scalar].abs().amax(0).clamp_min(tiny) + slope_scale = first[:, scalar].abs().amax(0) + + # === Step 1. Leading singular subspace of the scaled channels === + samples = ( + torch.cat( + [values[:, scalar], first[:, scalar] * _LOW_RANK_SLOPE_LENGTH], + dim=0, + ) + / value_scale + ) + if samples.shape[0] <= profile.basis or not bool(torch.isfinite(samples).all()): + return empty() + _, _, right = torch.linalg.svd(samples, full_matrices=False) + # Rows of ``right`` follow decreasing singular values; the flip stores the + # basis in increasing order. + subspace = right[: profile.basis].flip(0).T + expansion = value_scale[:, None] * subspace + + def project(knot_data: torch.Tensor) -> torch.Tensor: + return (knot_data[:, scalar] / value_scale) @ subspace + + basis_values = project(values) + basis_first = project(first) + basis_second = project(second) + + # === Step 2. Adopt the basis only if it reproduces every channel === + value_error = (basis_values @ expansion.T - values[:, scalar]).abs().amax(0) + slope_error = (basis_first @ expansion.T - first[:, scalar]).abs().amax(0) + if not bool( + (value_error <= _LOW_RANK_TOLERANCE * value_scale).all() + and (slope_error <= _LOW_RANK_TOLERANCE * slope_scale).all() + ): + return empty() + + # === Step 3. Tabulate the degree-one channels, the basis and the modes === + def columns(knot_data: torch.Tensor, basis: torch.Tensor) -> torch.Tensor: + return torch.cat( + [knot_data[:, :split], basis, knot_data[:, profile.channels :]], + dim=-1, + ) + + table = _quintic_coefficients( + columns(values, basis_values), + columns(first, basis_first), + columns(second, basis_second), + float(stride), + ) + return ( + table.to(torch.float32).contiguous(), + expansion.T.to(torch.float32).contiguous(), + ) def build_compression_artifacts( @@ -663,9 +932,10 @@ def build_compression_artifacts( Returns ------- dict[str, torch.Tensor] - Radial spline coefficients, metadata, ordered FiLM, mixing and type - caches, padded readout projections, sparse angular couplings, and the - output calibration. + Radial spline coefficients, the low-rank table of the scalar-only + channels and its expansion (empty when no basis is adopted), metadata, + ordered FiLM, mixing and type caches, padded readout projections, + sparse angular couplings, and the output calibration. Raises ------ @@ -673,34 +943,19 @@ def build_compression_artifacts( If the descriptor excludes type pairs, or is not an fp32 model with a compiled specialization. """ - if getattr(descriptor, "exclude_types", None): - raise ValueError( - "DPA4C compressed CUDA has no type-exclusion branch, so a " - "descriptor with a non-empty `exclude_types` cannot reach the " - "fused path. Drop the exclusions, or deploy the uncompressed " - "graph archive." - ) + reason = compression_ineligibility_reason(descriptor) + if reason is not None: + raise ValueError(reason) descriptor_state = descriptor.stddev - if descriptor_state.dtype != torch.float32: - raise ValueError( - "DPA4C compressed CUDA requires descriptor precision `float32`, " - f"got {descriptor_state.dtype}" - ) - if not mega_eligible(descriptor): - raise ValueError( - "DPA4C compressed CUDA supports " - f"channels {SUPPORTED_CHANNELS}, lmax {SUPPORTED_LMAX} and " - f"radial_modes {SUPPORTED_RADIAL_MODES}, got " - f"channels={descriptor.channels}, lmax={descriptor.lmax}, " - f"radial_modes={descriptor.radial_modes}" - ) device = descriptor_state.device profile = descriptor_profile( descriptor.channels, descriptor.lmax, descriptor.spin is not None, ) - table, info = build_radial_table(descriptor, stride) + knots = _radial_knots(descriptor, stride) + table, info = build_radial_table(descriptor, stride, knots=knots) + basis_table, basis_matrix = build_low_rank_table(descriptor, stride, knots=knots) with torch.no_grad(): readout_matrices = _build_readout_matrices(descriptor, profile, device) @@ -727,6 +982,8 @@ def build_compression_artifacts( to_device = {"device": device} return { "data": table, + "basis_table": basis_table, + "basis_matrix": basis_matrix, "info": info, "spin_type": _build_spin_type_cache(descriptor, device), "readout_matrices": readout_matrices, @@ -1205,6 +1462,7 @@ def _reference_descriptor( spin: torch.Tensor | None = None, spin_pair: torch.Tensor | None = None, spin_type: torch.Tensor | None = None, + edge_spin: torch.Tensor | None = None, ) -> torch.Tensor: """Reference implementation of the compressed DPA4C descriptor. @@ -1216,7 +1474,11 @@ def _reference_descriptor( ``spin`` with shape ``(N, 3)``, the ordered scale and shift cache ``spin_pair`` with shape ``((T + 1) ** 2, C_s, 2)``, and the per-type scalars ``spin_type`` with shape ``(T + 1, 4)``. A spin-free descriptor - omits it and reproduces the geometric descriptor exactly. + omits it and reproduces the geometric descriptor exactly. The optional + ``edge_spin`` with shape ``(E, 3)`` holds the moment of the source of + every edge, which the neighbor families read in place of + ``spin[source]``; with equal values the descriptor is the same, and + the cotangent of the neighbor moments is resolved per edge. """ del destination_order, destination_row_ptr, canonical, coupling_meta compute = edge_vec.to(torch.float32) @@ -1367,6 +1629,7 @@ def _reference_descriptor( if has_spin: magnitude, coordination, spin_vector, spin_tensor = _reference_spin_moments( spin, + spin[source] if edge_spin is None else edge_spin, spin_pair, spin_type, atype, @@ -1415,6 +1678,7 @@ def _reference_descriptor( def _reference_spin_moments( spin: torch.Tensor, + edge_spin: torch.Tensor, spin_pair: torch.Tensor, spin_type: torch.Tensor, atype: torch.Tensor, @@ -1454,6 +1718,8 @@ def _reference_spin_moments( ---------- spin Raw per-node magnetic moments with shape ``(N, 3)``. + edge_spin + Raw moment of the source node of every edge with shape ``(E, 3)``. spin_pair Ordered spin scale and shift with shape ``((T + 1) ** 2, C_s, 2)``. spin_type @@ -1495,7 +1761,7 @@ def _reference_spin_moments( spin_channels = int(spin_pair.shape[1]) weights = spin_type[atype] conditioned = spin.to(envelope.dtype) * weights[:, 0:1] - neighbor = conditioned[source] + neighbor = edge_spin.to(envelope.dtype) * weights[source, 0:1] film = spin_pair[pair_index] weight = (radial[:, :spin_channels] * film[..., 0] + film[..., 1]) * ( envelope * envelope @@ -1659,8 +1925,23 @@ def _reference_pair_energy( return energy.index_add(0, destination, 0.5 * series / radius * edge_mask)[:, None] -def _reference_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: +def _portable_inputs(args: tuple) -> tuple: + """Drop the low-rank table and its expansion from the operator inputs. + + ``args`` are the inputs that follow ``edge_vec``. Only the device scans + read the low-rank table, those of a system of one node type and the + slotted scans of a system of a few; the portable evaluation reads the + radial table itself, so it consumes the remaining inputs. + """ + return (*args[:6], *args[8:]) + + +def _reference_forward( + edge_vec: torch.Tensor, + *args: Any, +) -> tuple[torch.Tensor, torch.Tensor]: """CPU custom-op implementation returning descriptor and opaque state.""" + args = (edge_vec, *_portable_inputs(args)) descriptor = _reference_descriptor( *args[:16], *args[19:], @@ -1690,6 +1971,8 @@ def _forward_fake( destination_row_ptr: torch.Tensor, atype: torch.Tensor, table: torch.Tensor, + basis_table: torch.Tensor, + basis_matrix: torch.Tensor, pair_film: torch.Tensor, pair_mixing: torch.Tensor, type_embedding: torch.Tensor, @@ -1719,6 +2002,8 @@ def _forward_fake( destination_order, destination_row_ptr, table, + basis_table, + basis_matrix, pair_film, pair_mixing, readout_matrices, @@ -1758,8 +2043,9 @@ def _forward_fake( #: Position of the pair table among the backward inputs that follow -#: ``edge_vec``: nineteen tensors, ``canonical``, ``lmax`` and seven geometry -#: scalars precede it, and the pair seed closes the list. +#: ``edge_vec`` once :func:`_portable_inputs` has dropped the low-rank table: +#: nineteen tensors, ``canonical``, ``lmax`` and seven geometry scalars precede +#: it, and the pair seed closes the list. _PAIR_TABLE_SLOT = 28 @@ -1770,6 +2056,7 @@ def _backward_fake( *args: Any, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: del descriptor_gradient, state + args = _portable_inputs(args) atype, spin, pair_table = args[4], args[15], args[_PAIR_TABLE_SLOT] has_spin = spin.ndim == 2 # Each absent output is allocated separately: the schema declares four @@ -1797,12 +2084,12 @@ def _reference_backward( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """CPU custom-op backward returning the coordinate and magnetic cotangents. - The device operator splits the magnetic cotangent because a - destination-major scan does not own the source of an edge, and the caller - closes it by reducing the per-edge part onto source nodes. The reference - differentiates the whole node axis at once, so it returns the complete - cotangent on the node axis and leaves the per-edge part at zero, which the - same reduction carries through unchanged. + The magnetic cotangent is split as in the device operator, which splits it + because a destination-major scan does not own the source of an edge: the + node axis carries the on-site part, and the per-edge part, the cotangent + of the source moment read by each edge, is what the caller reduces onto + the source nodes. The reference differentiates the source moments as + their own per-edge leaf to obtain it. A non-empty pair table adds the analytical pair potential: its energy, weighted by the per-node cotangent ``pair_seed``, joins the differentiated @@ -1810,11 +2097,17 @@ def _reference_backward( and the node energies of the term are returned as the fourth result. """ del state + args = _portable_inputs(args) spin = args[15] pair_table, pair_seed = args[_PAIR_TABLE_SLOT:] has_spin = spin.ndim == 2 value = edge_vec.detach().clone().requires_grad_(True) moment = spin.detach().clone().requires_grad_(has_spin) + edge_moment = ( + spin.detach()[args[0][0].to(torch.long)].requires_grad_(True) + if has_spin + else None + ) pair_energy = edge_vec.new_empty((0,), dtype=torch.float64) with torch.enable_grad(): descriptor = _reference_descriptor( @@ -1824,6 +2117,7 @@ def _reference_backward( spin=moment, spin_pair=args[16], spin_type=args[17], + edge_spin=edge_moment, ) objective = (descriptor * descriptor_gradient.to(descriptor.dtype)).sum() if pair_table.numel(): @@ -1838,12 +2132,12 @@ def _reference_backward( objective = objective + (pair_energy[:, 0] * pair_seed.reshape(-1)).sum() gradients = torch.autograd.grad( objective, - (value, moment) if has_spin else (value,), + (value, moment, edge_moment) if has_spin else (value,), ) return ( gradients[0].to(edge_vec.dtype), gradients[1] if has_spin else edge_vec.new_empty((0,), dtype=torch.float32), - torch.zeros_like(edge_vec) + gradients[2].to(edge_vec.dtype) if has_spin else edge_vec.new_empty((0,), dtype=torch.float32), pair_energy.detach(), @@ -1853,13 +2147,17 @@ def _reference_backward( #: Position of the per-node magnetic moment among the operator inputs, and its #: position among the tensors :func:`_setup_context` saves, which lead with the #: opaque state. -_SPIN_INPUT_SLOT = 16 +_SPIN_INPUT_SLOT = 18 _SPIN_SAVED_SLOT = 1 + _SPIN_INPUT_SLOT +#: Tensor inputs of the descriptor operator: ``edge_vec`` and the twenty-one +#: tensors that follow it. The scalars come after them. +_TENSOR_INPUT_COUNT = 22 + def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: - ctx.save_for_backward(output[1], *inputs[:20]) - ctx.scalars = inputs[20:] + ctx.save_for_backward(output[1], *inputs[:_TENSOR_INPUT_COUNT]) + ctx.scalars = inputs[_TENSOR_INPUT_COUNT:] ctx.mark_non_differentiable(output[1]) ctx.set_materialize_grads(False) @@ -1902,7 +2200,7 @@ def _backward( no_pair, no_pair, )[0] - return (edge_gradient,) + (None,) * 28 + return (edge_gradient,) + (None,) * (_TENSOR_INPUT_COUNT - 1 + len(ctx.scalars)) @cache @@ -1980,13 +2278,15 @@ def compressed_operator_arguments( Returns ------- tuple - Radial table, ordered caches, readout projections, coupling tables, - output calibration, the native spin block, and the per-element length - scales of the bridging window. + Radial table, low-rank table and its expansion, ordered caches, + readout projections, coupling tables, output calibration, the native + spin block, and the per-element length scales of the bridging window. """ empty = descriptor.compress_spin_type.new_empty(0) return ( descriptor.compress_data, + descriptor.compress_basis_table, + descriptor.compress_basis_matrix, descriptor.compress_pair_film, descriptor.compress_pair_mixing, descriptor.compress_type_embedding, diff --git a/deepmd/pt_expt/kernels/fitting_routes.py b/deepmd/pt_expt/kernels/fitting_routes.py new file mode 100644 index 0000000000..354668107b --- /dev/null +++ b/deepmd/pt_expt/kernels/fitting_routes.py @@ -0,0 +1,795 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Freeze-time measurement of the fitting route of compressed DPA4C. + +The fused energy-force operator of compressed DPA4C evaluates its fitting +network over runs of nodes, either through the fused kernel, in one of the +shapes compiled for the hidden width, or through the per-layer kernels. Every +route evaluates the same contractions in the same order and so returns the +same bits; the routes differ only in speed. At a common clock a route takes +the time of the waves that its grids take on the device, which depends on how +full the last of them is, while a board at its power limit clocks each route by +the power that the route itself draws. The operator derives a route from rules +when it is given no table. At freeze time, on the deployment device, +:func:`tune_fitting_routes` measures instead which route serves each run size +and records the result as a fitting-route table, which names a route only +where the measurement departs from the rules and leaves them every other run +size, so that the frozen program runs the code of an untuned one wherever the +two agree. The table is passed to the operator as a literal of the frozen +program (:func:`fitting_routes`), which falls back to the rules on a device of +another fingerprint. + +The run sizes are divided into bands between which no compared route changes +its number of waves and no rule changes its route +(``graph_fitting_route_bands``). A route is fastest over a whole band only if +the fill of its last wave does not decide; the measurement therefore times both +ends of every band, and where the two ends disagree it bisects the band for the +run size at which the winner changes. Each sample times the compared routes on +the canonical graph of a periodic bulk system of exactly the run size of +atoms, as replays of a captured CUDA graph of the whole operator, which is how +LAMMPS evaluates it. + +A sample times the routes in two ways. Within each capture of the operator the +replays of the routes alternate, so that every route runs at the clocks of a +lightly loaded board, and a route other than that of the rules is a candidate +only where its median gain over the rules exceeds :data:`TIE` and the +half-range of its per-capture gains, half the difference between the largest +and the smallest of these gains; a deployment captures the operator once and +sees the gain of one capture. At a common clock the fitting network takes the +same time whatever the node types of the system, so the alternating replays run +the node types of the sample alone. A train is a run of consecutive replays of +one route, timed after a lead-in of further replays of that route; trains time +the routes at the clocks of a board held at its power limit, where a route that +draws more power can lose what it wins at a common clock. The loss depends on +the node types of the system, its composition: the descriptor may run other +scans for one node type than for several, and so sets the power and the length +of the step around the fitting network. On an L20 at its 350 W limit, the +per-layer kernels of a trained Air model at 12,320 nodes gain 1.1% on trains +with every node carbon and lose 1.4% with three node types. A candidate +therefore takes a run size only where it loses by no more than :data:`TIE` on +the trains of two compositions: every node of one type, and the node types of +the sample. Where a deployment leaves the device idle part of the time, the +gain of a route lies between its gains in the alternating replays and on the +trains, as on a V100 at its 300 W limit, where the gain on trains rises towards +that of the alternating replays as the device idles more. On these compositions +a route that the measurement selects is therefore not slower than the rules at +a load in between, beyond the tie and the noise of the measurements. + +A latent-charge long-range module evaluated by the fused long-range operator +runs its charge network through the same routes, by a table of its own. The +measurement times that network in place of the energy fitting of the fused +operator without a module, which runs the same kernels on runs of the same +sizes; a charge network of the key of the energy fitting shares its table. +""" + +from __future__ import ( + annotations, +) + +import dataclasses +import logging +import math +import statistics +import time +from typing import ( + Any, +) + +import torch + +from deepmd.pt_expt.kernels.autotune import ( + register_autotuner, +) + +log = logging.getLogger(__name__) + +#: First word of a fitting-route table (graph_fitting.cu). +ROUTE_TABLE_MAGIC = 0x52545431 +#: Route of the per-layer kernels (graph_ops.h). +PER_LAYER = -1 +#: Choice of a table that leaves the run to the rules of the operator +#: (graph_ops.h). +RULES = 0 +#: Relative gain within which a route counts as equally fast as that of the +#: rules. A route other than that of the rules is a candidate where its median +#: gain over the rules in the alternating replays of a sample exceeds its bound +#: (:func:`_bound`), the tie or a larger half-range of its per-capture gains. A +#: candidate replaces the route of the rules only if it loses to that route by +#: no more than the tie on the trains of one route at a time. +TIE = 0.01 +#: Atoms per ų of the synthetic systems: diamond, the density of the benchmark +#: systems. +DENSITY = 0.176 +#: Seed of the synthetic systems. The system of a run size is drawn from this +#: seed and the size, so that it does not depend on the order of the samples. +SEED = 20261018 +#: Largest run size measured for an operator that evaluates the whole node +#: axis at once (a node tile of zero), beyond which the choice of the last +#: interval extends; an operator of node tiles is measured up to its tile. +MAX_RUN_NODES = 131_072 +#: Run sizes over which every route keeps the grids of its kernels. +BLOCK = 32 +#: Node types of the synthetic systems, drawn at random up to the type count of +#: the model, as a system of a few elements. +SYNTHETIC_TYPES = 3 +#: Captures of the operator per route in the stages of one sample, and the +#: timed replays of every capture. The speed of a captured graph depends on +#: where its memory lies: at large runs the medians of different captures of +#: one route scatter by several percent, more than the replays of one capture +#: do, so a sample decides on the gains of several captures. The first replays of a +#: capture understate the differences between routes, so every capture takes +#: untimed warm-up replays first. +STAGE_CAPTURES = (1, 2, 4) +STAGE_REPLAYS = 3 +WARMUP_REPLAYS = 2 +#: Margins of the stages but the last. A sample decides after such a stage +#: where the median gain of every route other than that of the rules lies +#: farther than the margin of the stage from its bound, and after the first +#: stage only where no route exceeds its bound, so that a departure from the +#: rules rests on at least three captures, those of the first two stages; the +#: last stage always decides. +STAGE_MARGINS = (0.04, 0.015) +#: Replays per capture of a route whose replay takes less than +#: SHORT_REPLAY_MS, which cost little. +SHORT_STAGE_REPLAYS = 9 +SHORT_REPLAY_MS = 2.0 +#: Least device time, in ms, of the replays of a route that lead in to its +#: timed train, and the device time of the chunks in which the lead-in is +#: timed. A board at its power limit settles the clock of its multiprocessors +#: on the power draw of a route within about 0.4 s of switching to it (an L20 +#: at its 350 W limit), so a train is timed only after half a second of its own +#: replays. +LEAD_IN_MS = 500.0 +LEAD_IN_CHUNK_MS = 50.0 +#: Least number of timed replays of a train, and the device time, in ms, that +#: they span at about the replay time of the lead-in. +TRAIN_REPLAYS = 9 +TRAIN_MS = 100.0 +#: Captures of the trains of each composition of a sample +#: (:meth:`_Bench._train_gain`). +TRAIN_CAPTURES = 2 + +#: Measured tables: (device fingerprint, network key) -> fitting-route table. +_ROUTE_TABLES: dict[tuple, tuple[int, ...]] = {} +#: Fingerprints of the CUDA devices measured in this process, by device index; +#: a table of the backend device is read while tracing, which calls no operator. +_FINGERPRINTS: dict[int, tuple[int, ...]] = {} + + +def network_key(fitting: Any) -> tuple | None: + """Return the key of a fitting network in the route tables. + + Parameters + ---------- + fitting : Any + The pt_expt fitting network: an energy fitting or the charge network + of a long-range module. + + Returns + ------- + tuple or None + Input width, hidden widths, identity residuals and activation code, or + ``None`` for a network the fused operator cannot evaluate. + """ + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_eligible, + fitting_operator_arguments, + ) + + if not fitting_eligible(fitting): + return None + network = fitting_operator_arguments(fitting) + widths = tuple(int(w.shape[1]) for w in network.weights) + return ( + int(network.weights[0].shape[0]), + widths, + tuple(int(r) for r in network.residuals), + int(network.activation), + ) + + +def _fingerprint(device: torch.device) -> tuple[int, ...]: + """Return the fingerprint that a table measured on ``device`` carries.""" + index = device.index or 0 + if index not in _FINGERPRINTS: + _FINGERPRINTS[index] = tuple( + int(v) for v in torch.ops.deepmd.graph_fitting_route_fingerprint(index) + ) + return _FINGERPRINTS[index] + + +def fitting_routes(fitting: Any) -> list[int]: + """Return the fitting-route table of a network for the backend device. + + Parameters + ---------- + fitting : Any + The pt_expt fitting network: an energy fitting or the charge network + of a long-range module. + + Returns + ------- + list[int] + The table that :func:`tune_fitting_routes` measured for the network on + the backend device in this process, or an empty list, with which the + operator takes the routes of its rules. + """ + if not _ROUTE_TABLES: + return [] + import deepmd.pt_expt.utils.env as env + + fingerprint = ( + _FINGERPRINTS.get(env.DEVICE.index or 0) if env.DEVICE.type == "cuda" else None + ) + key = network_key(fitting) + if fingerprint is None or key is None: + return [] + return list(_ROUTE_TABLES.get((fingerprint, key), ())) + + +def _bulk_graph(rcut: float, node_count: int, device: torch.device) -> Any: + """Return the canonical graph of a periodic bulk system of ``node_count`` atoms. + + The atoms occupy ``node_count`` sites, drawn at random, of a simple cubic + lattice of ``ceil(node_count ** (1/3))`` sites per side in a periodic box + of the density :data:`DENSITY`, each displaced by up to 0.1 Å. The density + is uniform, every atom has the full neighbourhood of the bulk, and the + nearest distances stay near the site spacing. The edges are all pairs + within ``rcut``, the periodic images of a box smaller than ``2 * rcut`` + included, an atom's own images among them: the ghost-free canonical graph + that LAMMPS gives ``node_count`` local atoms. They are sorted by bond + length, as the LAMMPS builder writes the rows. + """ + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_pairs, + ) + from deepmd.pt_expt.utils.vesin_search import ( + VesinNeighborSearch, + ) + + generator = torch.Generator(device=device).manual_seed(SEED + node_count) + side = math.ceil(node_count ** (1.0 / 3.0)) + length = (node_count / DENSITY) ** (1.0 / 3.0) + sites = torch.randperm( + side**3, dtype=torch.int64, device=device, generator=generator + )[:node_count] + cells = torch.stack((sites // side**2, sites // side % side, sites % side), dim=-1) + positions = (cells.to(torch.float64) + 0.5) * (length / side) + positions += 0.1 * ( + 2.0 + * torch.rand( + positions.shape, dtype=torch.float64, device=device, generator=generator + ) + - 1.0 + ) + box = torch.eye(3, dtype=torch.float64, device=device) * length + center, source, vector = VesinNeighborSearch().compute( + positions, box, rcut, return_vectors=True + ) + order = torch.argsort(torch.linalg.vector_norm(vector, dim=-1), stable=True) + n_node = torch.tensor([node_count], dtype=torch.int64, device=device) + graph = canonical_graph_from_pairs( + center[order], source[order], vector[order], n_node, node_count + ) + # The operator reads the source CSR only for the per-atom virial, which + # LAMMPS declines; it takes two-entry placeholders in its place. + return dataclasses.replace( + graph, + source_row_ptr=graph.source_row_ptr.new_zeros(2), + source_order=graph.source_order.new_zeros(2), + ) + + +class _Bench: + """Times the fused operator of one model on synthetic bulk systems. + + The operator is that without a long-range module, with the descriptor of + the model and the network ``fitting`` with its atomic bias ``atom_bias``, + the energy fitting or a charge network, in place of the energy fitting. + The bench keeps four counters, which :func:`_tune_network` logs: + + - ``samples``, the samples of run sizes that it times, those of the + bisections of bands included; + - ``checked``, the samples that it times on trains; + - ``overturned``, the samples in which the trains refuse the candidate of + the largest gain in the alternating replays; a sample of two routes + besides that of the rules, which runs no alternating replays, never + counts; + - ``replays``, the replays of the alternating rounds of the samples and + those of the trains, their lead-ins included. + """ + + def __init__( + self, + model: Any, + device: torch.device, + largest: int, + fitting: Any, + atom_bias: torch.Tensor, + ) -> None: + import copy + + learned, pair_potential = model.atomic_model.fused_decomposition() + self.descriptor = copy.deepcopy(learned.descriptor).to(device).eval() + self.fitting = copy.deepcopy(fitting).to(device).eval() + self.atom_bias = atom_bias.to(device) + self.pair_table = ( + None if pair_potential is None else pair_potential.pair_table.to(device) + ) + self.device = device + self.rcut = float(self.descriptor.get_rcut()) + generator = torch.Generator(device=device).manual_seed(SEED) + self.types = min(SYNTHETIC_TYPES, int(self.descriptor.get_ntypes())) + self.atype = torch.randint( + 0, self.types, (largest,), device=device, generator=generator + ) + self.one_type = torch.zeros_like(self.atype) + self.spin = None + if getattr(self.descriptor, "spin", None) is not None: + moments = torch.randn(largest, 3, device=device, generator=generator) + self.spin = moments / torch.linalg.vector_norm( + moments, dim=-1, keepdim=True + ) + self.fingerprint = _fingerprint(device) + self.samples = 0 + self.checked = 0 + self.overturned = 0 + self.replays = 0 + + def single(self, choice: int) -> list[int]: + """Return the table that names ``choice`` for every run size.""" + return [ROUTE_TABLE_MAGIC, *self.fingerprint, 1, choice] + + def _inputs(self, node_count: int) -> tuple[Any, torch.Tensor, Any]: + """Return the graph, the types and the moments of one run size.""" + graph = _bulk_graph(self.rcut, node_count, self.device) + spin = None if self.spin is None else self.spin[:node_count] + return graph, self.atype[:node_count], spin + + def _compositions(self, inputs: tuple) -> list[tuple]: + """Return the compositions of the trains of a sample. + + The trains run the graph of the sample with every node of the first + node type and with the node types of the sample; for a model of one + node type the two coincide, and the trains run it once. + """ + graph, atype, spin = inputs + one_type = (graph, self.one_type[: atype.shape[0]], spin) + return [one_type] if self.types == 1 else [one_type, inputs] + + def _capture( + self, inputs: tuple, choices: list[int], rotation: int + ) -> dict[int, Any]: + """Return a captured CUDA graph of the operator on ``inputs`` per route. + + The routes are captured in the order of ``choices`` rotated by + ``rotation``, so that over the captures of a sample no route keeps one + place in the shared memory pool. + """ + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + dpa4c_canonical_compress_energy_force, + ) + + graph, atype, spin = inputs + stream = torch.cuda.Stream(self.device) + graphs: dict[int, Any] = {} + # The graphs of one sample share one memory pool; they replay one at a + # time on one stream, and nothing reads their outputs. + pool = None + shift = rotation % len(choices) + for choice in choices[shift:] + choices[:shift]: + table = self.single(choice) + + def call(table: list[int] = table) -> Any: + return dpa4c_canonical_compress_energy_force( + self.descriptor, + self.fitting, + graph, + atype, + self.atom_bias, + spin, + self.pair_table, + fitting_routes=table, + ) + + stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(stream): + call() + torch.cuda.current_stream(self.device).wait_stream(stream) + captured = torch.cuda.CUDAGraph() + with torch.cuda.graph(captured, pool=pool, stream=stream): + call() + pool = captured.pool() + graphs[choice] = captured + for captured in graphs.values(): + for _ in range(WARMUP_REPLAYS): + captured.replay() + torch.cuda.synchronize(self.device) + return graphs + + def _replay( + self, graphs: dict[int, Any], replays: int, times: dict[int, list[float]] + ) -> None: + """Append the times of ``replays`` alternating replays of every route.""" + choices = list(graphs) + events = [] + for replay in range(replays): + shift = replay % len(choices) + for choice in choices[shift:] + choices[:shift]: + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + graphs[choice].replay() + end.record() + events.append((choice, begin, end)) + torch.cuda.synchronize(self.device) + for choice, begin, end in events: + times[choice].append(begin.elapsed_time(end)) + self.replays += replays * len(choices) + + def winner(self, node_count: int, rule: int, choices: list[int]) -> int: + """Return the fastest route on one run size. + + Within one capture of the operator the replays of the routes + alternate, so every route replays at the same clocks of the device, + and the gain of a route over the route of the rules is the ratio of + their median replay times, less one. A route is a candidate where its + median gain over the captures of the sample exceeds its bound + (:func:`_bound`); a deployment captures the operator once and replays + that capture, so a gain within the half-range of the per-capture gains + is not one that it reliably sees. The sample takes further captures + while the gain of any route lies near its bound, and a departure from + the rules rests on at least three captures (:data:`STAGE_MARGINS`). + The candidates are checked on trains of one route at a time in the + order of their gains, and the first that holds there takes the run + size (:meth:`_departure`). + + The bisection of a band compares the two routes that won at its ends. + Where neither of them is that of the rules, ``choices`` holds these two + alone, and the trains decide between them (:meth:`_faster_on_trains`); + where one of them is, the sample runs as above. + """ + if len(choices) == 1: + return choices[0] + self.samples += 1 + inputs = self._inputs(node_count) + if rule not in choices: + first, second = choices + return self._faster_on_trains(inputs, first, second) + # One untimed round brings the clocks of the device to the load of the + # sample, and sets the replays of its captures. + graphs = self._capture(inputs, choices, 0) + short = self._replay_ms(graphs[choices[0]]) < SHORT_REPLAY_MS + replays = SHORT_STAGE_REPLAYS if short else STAGE_REPLAYS + self._replay(graphs, replays, {choice: [] for choice in choices}) + del graphs + rivals = [choice for choice in choices if choice != rule] + gains: dict[int, list[float]] = {choice: [] for choice in rivals} + captures = 0 + for stage, count in enumerate(STAGE_CAPTURES): + for _ in range(count): + graphs = self._capture(inputs, choices, captures) + captures += 1 + times: dict[int, list[float]] = {choice: [] for choice in choices} + self._replay(graphs, replays, times) + del graphs + medians = {choice: statistics.median(v) for choice, v in times.items()} + for choice in rivals: + gains[choice].append(medians[rule] / medians[choice] - 1.0) + margins = [ + statistics.median(gains[choice]) - _bound(gains[choice]) + for choice in rivals + ] + settled = stage > 0 or max(margins) <= 0.0 + if stage == len(STAGE_MARGINS) or ( + settled + and all(abs(margin) > STAGE_MARGINS[stage] for margin in margins) + ): + break + candidates = sorted( + ( + choice + for choice, margin in zip(rivals, margins, strict=True) + if margin > 0.0 + ), + key=lambda choice: statistics.median(gains[choice]), + reverse=True, + ) + return self._departure(inputs, rule, candidates) + + def _departure(self, inputs: tuple, rule: int, candidates: list[int]) -> int: + """Return the first of ``candidates`` that holds on trains, or the rules. + + A candidate holds where its median gain over the route of the rules on + the trains is at least ``-TIE`` for every composition of + :meth:`_compositions`: within the tie, two routes count as equally fast, + on trains as in the alternating replays. The check of a candidate stops + at the first composition on which it loses by more. A sample counts as + overturned where the trains refuse the candidate of the largest gain, + whether a later candidate or the route of the rules then takes the run + size. + """ + if not candidates: + return rule + self.checked += 1 + for candidate in candidates: + if all( + self._train_gain(nodes, rule, candidate) >= -TIE + for nodes in self._compositions(inputs) + ): + self.overturned += candidate != candidates[0] + return candidate + self.overturned += 1 + return rule + + def _faster_on_trains(self, inputs: tuple, first: int, second: int) -> int: + """Return the faster of two routes on trains of one route at a time. + + The second route is faster where the mean, over the compositions of + :meth:`_compositions`, of its median gains over the first on the trains + is positive, and the first route otherwise. + """ + self.checked += 1 + gains = [ + self._train_gain(nodes, first, second) + for nodes in self._compositions(inputs) + ] + return second if statistics.mean(gains) > 0.0 else first + + def _train_gain(self, inputs: tuple, incumbent: int, rival: int) -> float: + """Return the median gain of ``rival`` over ``incumbent`` on trains. + + Each capture times a train of each route (:meth:`_train_ms`); the + captures alternate the order of the two routes, so that a drift of the + clocks cancels in the median gain. + """ + gains = [] + for capture in range(TRAIN_CAPTURES): + graphs = self._capture(inputs, [incumbent, rival], capture) + order = (incumbent, rival) if capture % 2 == 0 else (rival, incumbent) + times = {choice: self._train_ms(graphs[choice]) for choice in order} + del graphs + gains.append(times[incumbent] / times[rival] - 1.0) + return statistics.median(gains) + + def _train_ms(self, graph: Any) -> float: + """Return the median replay time, in ms, of a train of one route. + + The timed replays follow a lead-in of at least :data:`LEAD_IN_MS` of + device time, after which a board at its power limit has settled its + clocks on the route. The lead-in accumulates device time in chunks of + about :data:`LEAD_IN_CHUNK_MS`, sized from the time of its first + replay. The train spans about :data:`TRAIN_MS` of device time at the + mean replay time of the lead-in, in at least :data:`TRAIN_REPLAYS` + replays. + """ + replay_ms = self._replay_ms(graph) + chunk = max(1, math.ceil(LEAD_IN_CHUNK_MS / replay_ms)) + lead_in_ms, lead_in = replay_ms, 1 + while lead_in_ms < LEAD_IN_MS: + lead_in_ms += self._replay_ms(graph, chunk) + lead_in += chunk + count = max(TRAIN_REPLAYS, math.ceil(TRAIN_MS * lead_in / lead_in_ms)) + events = [] + for _ in range(count): + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + graph.replay() + end.record() + events.append((begin, end)) + torch.cuda.synchronize(self.device) + self.replays += lead_in + count + return statistics.median([begin.elapsed_time(end) for begin, end in events]) + + def _replay_ms(self, graph: Any, count: int = 1) -> float: + """Return the time, in ms, of ``count`` consecutive replays of ``graph``.""" + start = torch.cuda.Event(enable_timing=True) + stop = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(count): + graph.replay() + stop.record() + stop.synchronize() + return start.elapsed_time(stop) + + +def _bound(gains: list[float]) -> float: + """Return the median gain that a route must exceed in the alternating replays. + + The bound is :data:`TIE`, or the half-range of the per-capture gains of + the route where that is larger. + """ + return max(TIE, 0.5 * (max(gains) - min(gains))) + + +def _choices(device: torch.device, width: int, node_count: int) -> list[int]: + return [ + PER_LAYER, + *( + int(k) + for k in torch.ops.deepmd.graph_fitting_route_choices( + device.index or 0, width, node_count + ) + ), + ] + + +def _measure( + bench: _Bench, ws: list[torch.Tensor], largest: int +) -> tuple[list[int], list[int]]: + """Return the inclusive upper bounds and the choices of the measured intervals. + + The intervals cover the run sizes up to ``largest``; the last one extends + beyond it. + """ + device = bench.device + width = int(ws[0].shape[1]) + din = int(ws[0].shape[0]) + bounds = [ + int(b) + for b in torch.ops.deepmd.graph_fitting_route_bands( + device.index or 0, width, din, largest + ) + ] + intervals: list[tuple[int, int, int]] = [] + low = 1 + for high in bounds: + rule = int(torch.ops.deepmd.graph_fitting_route_rule(ws, high)) + choices = _choices(device, width, high) + if rule not in choices: + choices.append(rule) + first = bench.winner(low, rule, choices) + last = first if high == low else bench.winner(high, rule, choices) + # Every route keeps the grids of its kernels over a block of run sizes + # (32 k, 32 k + 32], since its panels and tiles hold 32 or 64 nodes; a + # band whose two ends lie in one block has one winner, the rules' if + # they take one of the two. + lo, hi = math.ceil(low / BLOCK), math.ceil(high / BLOCK) + if first != last and lo == hi: + first = last = rule if rule in (first, last) else last + if first == last: + intervals.append((high, first, rule)) + else: + # The winner changes inside the band: bisect over the blocks for + # the last block that the first winner serves. + while hi - lo > 1: + middle = (lo + hi) // 2 + if bench.winner(BLOCK * middle, rule, [first, last]) == first: + lo = middle + else: + hi = middle + intervals.append((BLOCK * lo, first, rule)) + intervals.append((high, last, rule)) + low = high + 1 + # The table records the rules for every interval that they route as + # measured, so that the operator evaluates them there as it does without a + # table; naming their route instead would skip that evaluation and its side + # effects, such as the loading of the fused kernels that the occupancy + # queries of the rules perform under lazy module loading. + merged: list[tuple[int, int]] = [] + for high, route, rule in intervals: + choice = RULES if route == rule else route + if merged and merged[-1][1] == choice: + merged[-1] = (high, choice) + else: + merged.append((high, choice)) + return [high for high, _ in merged[:-1]], [choice for _, choice in merged] + + +def _measured_networks(model: torch.nn.Module) -> list[tuple[Any, torch.Tensor]]: + """Return the networks whose routes the fused operators of a model take. + + They are the energy fitting with its combined atomic bias and, for a + long-range module that the fused long-range operator evaluates, its + charge network with its per-type bias. + """ + from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + fused_long_range_eligible, + ) + from deepmd.pt_expt.model.make_model import ( + fused_atom_bias, + ) + + learned, _ = model.atomic_model.fused_decomposition() + networks = [(learned.fitting_net, fused_atom_bias(model.atomic_model, learned))] + if learned.has_long_range() and fused_long_range_eligible(learned.long_range): + charge_net = learned.long_range.charge_net + networks.append((charge_net, charge_net.bias_atom_e[:, 0].detach())) + return networks + + +def _tune_network( + model: torch.nn.Module, fitting: Any, atom_bias: torch.Tensor, device: torch.device +) -> None: + """Measure the routes of one network of ``model`` on ``device``. + + A no-op for a network and device whose table is already measured in this + process, and for a width of which no fused shape runs on the device. + """ + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_operator_arguments, + node_tile, + ) + + key = network_key(fitting) + if key is None: + return + fingerprint = _fingerprint(device) + if (fingerprint, key) in _ROUTE_TABLES: + return + width = key[1][0] + # The frozen operator evaluates runs of at most its node tile, which the + # freeze reads as this process does. + largest = node_tile() or MAX_RUN_NODES + if not any( + torch.ops.deepmd.graph_fitting_route_choices(device.index or 0, width, nodes) + for nodes in (1, largest) + ): + # No fused shape of the width runs on the device, so the per-layer + # kernels serve every run, as the rules route them. + return + start = time.perf_counter() + # The neighbor search of the samples and the events that time their + # replays act on the current device, so the measurement makes the device + # of the freeze current; the search launched on another device than that + # of its positions would read them out of bounds. + with torch.cuda.device(device), torch.no_grad(): + bench = _Bench(model, device, largest, fitting, atom_bias) + network = fitting_operator_arguments(bench.fitting) + bounds, choices = _measure(bench, list(network.weights), largest) + table = (ROUTE_TABLE_MAGIC, *fingerprint, len(choices), *bounds, *choices) + _ROUTE_TABLES[(fingerprint, key)] = table + log.info( + "Measured the fitting routes of hidden width %d on %s in %.1f s " + "(%d samples, %d of them checked on trains of one route and %d of " + "these overturned, %d replays): %d intervals, bounds %s, choices %s", + width, + torch.cuda.get_device_name(device), + time.perf_counter() - start, + bench.samples, + bench.checked, + bench.overturned, + bench.replays, + len(choices), + bounds, + choices, + ) + + +def tune_fitting_routes( + model: torch.nn.Module, level: int, device: torch.device +) -> None: + """Measure the fitting routes of a compressed DPA4C model on ``device``. + + The routes of the energy fitting and, for a long-range module that the + fused long-range operator evaluates, of its charge network are measured + one network at a time. A no-op for a model that does not take the fused + energy-force operator of compressed DPA4C. + + Parameters + ---------- + model : torch.nn.Module + The model about to be frozen. + level : int + The ``DP_TRITON_INFER`` level, which the measurement does not read. + device : torch.device + The CUDA device the frozen model is compiled for. + """ + del level + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + canonical_model_eligible, + op_available, + ) + + if not op_available() or not canonical_model_eligible(model): + return + for fitting, atom_bias in _measured_networks(model): + _tune_network(model, fitting, atom_bias, device) + + +register_autotuner( + "dpa4c_fitting_routes", + tune_fitting_routes, + min_level=0, + lower_kinds=("dpa4c_canonical",), +) diff --git a/deepmd/pt_expt/kernels/graph_fitting.py b/deepmd/pt_expt/kernels/graph_fitting.py index 33c7209b59..c9074fd69d 100644 --- a/deepmd/pt_expt/kernels/graph_fitting.py +++ b/deepmd/pt_expt/kernels/graph_fitting.py @@ -2,13 +2,17 @@ """Bindings for the fused energy fitting operator of the graph lower. The operator ``deepmd::graph_fitting`` evaluates the whole energy fitting -network on the flat node axis -- GEMMs with the bias / activation / residual -epilogues fused into the surrounding elementwise pass -- and returns the -per-atom energy in fp64. The registered backward chains the layer dgrads from -the saved pre-activations, exposing the descriptor gradient that the force / -virial assembly differentiates through. The CUDA kernel is +network on the flat node axis -- one fp32 SIMT GEMM kernel per layer, with the +bias, activation, residual and energy head fused into its epilogue -- and +returns the per-atom energy in fp64. The registered backward chains the layer +dgrads from the saved pre-activations, exposing the descriptor gradient that +the force / virial assembly differentiates through. +``deepmd::graph_fitting_energy_gradient`` +evaluates the energy and the descriptor cotangent of an inference step +together, in one kernel for the whole network where that kernel serves the +run, with results bitwise those of the per-layer kernels. The CUDA kernels are ``source/op/pt/graph_fitting.cu`` and the CPU kernel -``source/op/pt/graph_fitting_cpu.cc``. +``source/op/pt/cpu/graph_fitting_cpu.cc``. The operator is descriptor-agnostic: any graph-lowered energy model whose fitting is a plain MLP over the flat node axis (see @@ -73,8 +77,8 @@ #: Nodes evaluated per run of a tiled inference pipeline. The tile bounds #: every node-scale allocation that a run retires, and at this size the extra -#: launches are a fraction of a percent of a saturated step while the GEMM -#: shape stays far above the width at which cuBLAS loses efficiency. +#: launches are a fraction of a percent of a saturated step while every GEMM +#: grid still spans many waves of the device. #: ``DP_NODE_TILE`` overrides it; zero evaluates the whole node axis at once. _DEFAULT_TILE = 131072 @@ -102,10 +106,9 @@ def fitting_eligible(fit: Any) -> bool: Requires a single mixed-types energy net with tanh / silu on every hidden layer, fp32 weights and biases, no layer timestep, a linear scalar head - without residual, float4-aligned hidden widths of at most 4096, and no - frame / atomic parameters, case embedding or type exclusion. Hidden - identity residuals are supported; width-doubling residuals use the - reference path. + without residual, float4-aligned hidden widths, and no frame / atomic + parameters, case embedding or type exclusion. Hidden identity residuals + are supported; width-doubling residuals use the reference path. A per-layer timestep is a supported configuration of the reference network but not of this operator. Nothing prevents it technically -- the @@ -155,12 +158,10 @@ def fitting_eligible(fit: Any) -> bool: layer.resnet and layer.w.shape[1] == 2 * layer.w.shape[0] for layer in hidden ): return False - # The elementwise epilogues map one float4 lane of a row to threadIdx.x, - # so a row must be non-empty and fit the 1024-thread block limit at four - # values per lane. + # The kernels address every hidden row, bias and head weight in aligned + # float4 quads, so each hidden width is a positive multiple of four. return all( - int(layer.w.shape[1]) % 4 == 0 and 0 < int(layer.w.shape[1]) <= 4096 - for layer in hidden + int(layer.w.shape[1]) % 4 == 0 and int(layer.w.shape[1]) > 0 for layer in hidden ) @@ -231,9 +232,7 @@ def fitting_operator_arguments(fit: Any) -> FittingArguments: residuals=[1 if layer.resnet else 0 for layer in hidden], head_weight=head.w.reshape(-1).contiguous(), head_bias=( - head.b.reshape(-1).to(torch.float32).contiguous() - if head.b is not None - else empty + head.b.to(torch.float32).contiguous() if head.b is not None else empty ), activation=ACT_CODES[str(hidden[0].activation_function).lower()], ) diff --git a/deepmd/pt_expt/kernels/long_range.py b/deepmd/pt_expt/kernels/long_range.py new file mode 100644 index 0000000000..33ab82a828 --- /dev/null +++ b/deepmd/pt_expt/kernels/long_range.py @@ -0,0 +1,480 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the operator of the latent-charge long range. + +The operator ``deepmd::latent_charge_sum`` evaluates, for the nodes of +complete frames, the potential that the charges of every other node of the +frame create at each node and the force between the charges. A non-periodic +frame is summed directly over its pairs without listing them, so that the +memory is linear in the node count; a periodic frame is summed over the images +of its cell by the smooth particle-mesh Ewald method, on a mesh sized by the +cell. The CUDA kernels are in ``source/op/pt/long_range``; the CPU kernel +registered here is the reference they are tested against. + +The operator serves inference. It applies the total-charge constraint of the +module, and returns the atomic energy, its derivative with respect to the +charges the charge network predicts, and the force and virial at fixed +charges; the caller adds the response of the charges by back-propagating the +charge gradient through the charge network. + +Given the MPI communicator of a system decomposed into domains, one per rank, +the CUDA kernel evaluates the charged nodes of its rank against those of +every rank; the CPU reference evaluates a single domain. +""" + +import math +from functools import ( + cache, +) +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + frame_id_from_n_node, + unordered_pair_index, +) +from deepmd.pt_expt.kernels.utils import ( + operator_available, +) +from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + SOGKernel, +) + +__all__ = [ + "MAX_CHANNELS", + "MAX_GAUSSIANS", + "MAX_ORDER", + "kernel_eligible", + "kernel_operator_arguments", + "latent_charge_sum", + "op_available", +] + +#: Charge channels and Gaussians the CUDA kernels are instantiated for. +MAX_CHANNELS = 8 +MAX_GAUSSIANS = 32 +#: The largest order of the B-splines of the mesh. +MAX_ORDER = 12 + +_LES = 0 +_SOG = 1 + + +def op_available() -> bool: + """Whether the backend device carries the operator. + + The CPU kernel is the reference of this module, which is registered here + so that the answer does not depend on earlier calls. + """ + # Importing the module loads the operator library. + import deepmd.pt.cxx_op # noqa: F401 + + ensure_registered() + return operator_available("latent_charge_sum") + + +def kernel_eligible(kernel: Any) -> bool: + """Whether the operator evaluates this interaction kernel.""" + if isinstance(kernel, LESKernel): + return True + return isinstance(kernel, SOGKernel) and kernel.amplitude.numel() <= MAX_GAUSSIANS + + +def kernel_operator_arguments(kernel: Any) -> tuple[int, torch.Tensor]: + """Convert an interaction kernel into the code and coefficients of the operator. + + Parameters + ---------- + kernel + A pt_expt interaction kernel accepted by :func:`kernel_eligible`. + + Returns + ------- + tuple[int, torch.Tensor] + The kernel code and its float64 coefficients: ``alpha`` for the erf + kernel; the amplitudes ``a_k`` and then the rates ``c_k = log2(e) / + (2 sigma_k^2)`` of ``a_k 2^(-c_k r^2)`` for the sum of Gaussians. + + Raises + ------ + ValueError + If the operator does not evaluate the kernel. + """ + if not kernel_eligible(kernel): + raise ValueError( + f"the long-range operator does not evaluate {type(kernel).__name__}" + ) + if isinstance(kernel, LESKernel): + coefficient = torch.exp(kernel.log_alpha.to(torch.float64)).reshape(1) + return _LES, coefficient.detach().contiguous() + rate = (0.5 / math.log(2.0)) * torch.exp(-2.0 * kernel.log_width.to(torch.float64)) + coefficient = torch.cat((kernel.amplitude.to(torch.float64), rate)) + return _SOG, coefficient.detach().contiguous() + + +def _kernel_and_slope( + r_sq: torch.Tensor, coefficient: torch.Tensor, kernel: int +) -> tuple[torch.Tensor, torch.Tensor]: + """K(r) and K'(r) / r at the squared distances ``r_sq``.""" + coefficient = coefficient.to(dtype=r_sq.dtype, device=r_sq.device) + if kernel == _LES: + alpha = coefficient[0] + # Coincident nodes, which carry no weight, are evaluated at a short + # distance, where the closed form stays finite. + r_sq = torch.clamp(r_sq, min=1e-12 / alpha**2) + r = torch.sqrt(r_sq) + value = torch.erf(alpha * r) / r + contact = (2.0 / math.sqrt(math.pi)) * alpha + slope = (contact * torch.exp(-(alpha**2) * r_sq) - value) / r_sq + return value, slope + amplitude, rate = coefficient.reshape(2, -1) + weight = torch.exp2(-r_sq[:, None] * rate) + return weight @ amplitude, weight @ (-2.0 * math.log(2.0) * rate * amplitude) + + +def _charges_of_frames( + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + channel: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The charges of the interactions in float64, shifted in the constrained channel. + + Returns the charges with shape ``(N, C)``, the weights of the charged + nodes and the frame of every node with shape ``(N,)``, and the number of + charged nodes of every frame, at least one, with shape ``(F,)``. + """ + n = charge.shape[0] + frame_count = n_node_per_frame.shape[0] + # Padding nodes belong to no frame and carry no charge. + frame = frame_id_from_n_node(n_node_per_frame, n_total=n) + real = torch.arange(n, dtype=torch.int64, device=charge.device) < torch.sum( + n_node_per_frame + ) + weight = (node_mask & real).to(torch.float64) + q = charge.to(torch.float64) * weight[:, None] + count = weight.new_zeros(frame_count).index_add(0, frame, weight) + count = torch.clamp(count, min=1.0) + if channel >= 0: + total = weight.new_zeros(frame_count).index_add(0, frame, q[:, channel]) + shift = (total - total_charge.to(torch.float64)) / count + q = q.clone() + q[:, channel] = q[:, channel] - shift[frame] * weight + return q, weight, frame, count + + +def _direct_reference( + coord: torch.Tensor, + q: torch.Tensor, + n_node_per_frame: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + prefactor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Potential, force and atomic virial over the explicit pairs of every frame.""" + n = coord.shape[0] + pairs = unordered_pair_index(n_node_per_frame) + src, dst = pairs[0], pairs[1] + d = coord[src] - coord[dst] + value, slope = _kernel_and_slope(torch.sum(d * d, dim=-1), coefficient, kernel) + potential = coord.new_zeros(q.shape) + potential = potential.index_add(0, src, q[dst] * value[:, None]) + potential = prefactor * potential.index_add(0, dst, q[src] * value[:, None]) + strength = torch.sum(q[src] * q[dst], dim=-1) * slope + pair_force = -strength[:, None] * d + force = torch.zeros_like(coord).index_add(0, src, pair_force) + force = prefactor * force.index_add(0, dst, -pair_force) + half = (-0.5 * strength)[:, None, None] * d[:, :, None] * d[:, None, :] + atom_virial = coord.new_zeros(n, 3, 3).index_add(0, src, half) + atom_virial = prefactor * atom_virial.index_add(0, dst, half).reshape(n, 9) + return potential, force, atom_virial + + +def _interaction_kernel(coefficient: torch.Tensor, kernel: int) -> Any: + """The dpmodel interaction kernel of the coefficients of the operator.""" + from deepmd.dpmodel.long_range.kernel import LESKernel as LESKernelDP + from deepmd.dpmodel.long_range.kernel import SOGKernel as SOGKernelDP + + values = coefficient.detach().to("cpu", torch.float64).numpy() + if kernel == _LES: + return LESKernelDP(alpha=float(values[0]), precision="float64") + amplitude, rate = values.reshape(2, -1) + width = (0.5 / (math.log(2.0) * rate)) ** 0.5 + return SOGKernelDP( + amplitude=amplitude.tolist(), width=width.tolist(), precision="float64" + ) + + +def _mesh_reference( + coord: torch.Tensor, + q: torch.Tensor, + weight: torch.Tensor, + n_node_per_frame: torch.Tensor, + box: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Potential, force and atomic virial on the mesh of every periodic frame. + + Every frame is evaluated on its own mesh by the evaluation of the training + path, and differentiated with respect to its charges, its positions and + its strain. The atomic virial takes ``F (x) r`` at every node and shares + the rest of the virial of the frame evenly over its charged nodes. + """ + from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, + ) + from deepmd.dpmodel.long_range.reciprocal import ( + mesh_energy, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_mesh_grid, + ) + + interaction = _interaction_kernel(coefficient, kernel) + potential = torch.zeros_like(q) + force = torch.zeros_like(coord) + atom_virial = coord.new_zeros(coord.shape[0], 3, 3) + begin = 0 + for f, length in enumerate(n_node_per_frame.tolist()): + end = begin + length + frame_q = q[begin:end].detach().requires_grad_(True) + frame_coord = coord[begin:end].detach().requires_grad_(True) + with torch.enable_grad(): + mesh = build_mesh_grid(box[f].reshape(1, 3, 3), spacing, frame_coord) + mesh.strain.requires_grad_(True) + frame_id = torch.zeros(length, dtype=torch.int64, device=coord.device) + energy = mesh_energy(frame_q, mesh, frame_id, interaction, order).sum() + energy = energy * (prefactor / COULOMB_CONSTANT) + d_q, d_coord, d_strain = torch.autograd.grad( + energy, (frame_q, frame_coord, mesh.strain) + ) + frame_force = -d_coord + moment = frame_force[:, :, None] * coord[begin:end, None, :] + frame_weight = weight[begin:end] + count = torch.clamp(torch.sum(frame_weight), min=1.0) + rest = (-d_strain[0] - torch.sum(moment, dim=0)) / count + potential[begin:end] = d_q + force[begin:end] = frame_force + atom_virial[begin:end] = moment + rest * frame_weight[:, None, None] + begin = end + return potential, force, atom_virial.reshape(-1, 9) + + +def _reference( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + box: torch.Tensor | None, + coefficient: torch.Tensor, + communicator: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the operator in float64, over the explicit pairs or the meshes.""" + if int(communicator.reshape(-1)[0]) != 0: + raise NotImplementedError( + "the CPU reference of latent_charge_sum evaluates a single domain; " + "a communicator of several ranks requires the CUDA kernel" + ) + frame_count = n_node_per_frame.shape[0] + q, weight, frame, count = _charges_of_frames( + charge, node_mask, n_node_per_frame, total_charge, channel + ) + has_cell = ( + torch.zeros(frame_count, dtype=torch.bool, device="cpu") + if box is None + else torch.any(box != 0.0, dim=-1) + ) + if bool(torch.any(has_cell)) != bool(torch.all(has_cell)): + raise ValueError( + "latent_charge_sum: every frame of a call is periodic, or none is" + ) + if bool(torch.any(has_cell)): + potential, force, atom_virial = _mesh_reference( + coord, + q, + weight, + n_node_per_frame, + box.to(coord.device), + coefficient, + kernel, + prefactor, + order, + spacing, + ) + else: + potential, force, atom_virial = _direct_reference( + coord, q, n_node_per_frame, coefficient, kernel, prefactor + ) + virial = coord.new_zeros(frame_count, 9).index_add(0, frame, atom_virial) + atom_energy = 0.5 * torch.sum(q * potential, dim=-1, keepdim=True) + gradient = potential * weight[:, None] + if channel >= 0: + mean = coord.new_zeros(frame_count).index_add(0, frame, gradient[:, channel]) + gradient = gradient.clone() + gradient[:, channel] = gradient[:, channel] - (mean / count)[frame] * weight + return ( + atom_energy, + gradient.to(charge.dtype), + force.to(charge.dtype), + atom_virial.to(charge.dtype), + virial.to(charge.dtype), + ) + + +def _fake( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + box: torch.Tensor | None, + coefficient: torch.Tensor, + communicator: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + n = coord.shape[0] + return ( + coord.new_empty(n, 1), + torch.empty_like(charge), + charge.new_empty(n, 3), + charge.new_empty(n, 9), + charge.new_empty(n_node_per_frame.shape[0], 9), + ) + + +@cache +def _register() -> None: + torch.library.register_fake("deepmd::latent_charge_sum")(_fake) + library = torch.library.Library("deepmd", "IMPL") + library.impl("latent_charge_sum", _reference, "CPU") + # The registration lives as long as the library object. + _register.library = library + + +def ensure_registered() -> None: + """Register the fake and CPU implementations when the schema is loaded.""" + if isinstance( + getattr(torch.ops.deepmd, "latent_charge_sum", None), + torch._ops.OpOverloadPacket, + ): + _register() + + +def latent_charge_sum( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + kernel: Any, + channel: int, + prefactor: float, + communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Energy, charge gradient, force and virial of the charges of complete frames. + + Parameters + ---------- + coord + Node positions with shape ``(N, 3)`` in Å, float64. The frames occupy + contiguous spans of the node axis; nodes past ``sum(n_node_per_frame)`` + are padding and join no frame. + charge + Raw node charges of the charge network with shape ``(N, C)``, + float32, with ``C`` at most :data:`MAX_CHANNELS`. + node_mask + Nodes that carry charges with shape ``(N,)``, bool; the others carry + none. + n_node_per_frame + Node count of every frame with shape ``(F,)``. + total_charge + Total charge of every frame with shape ``(F,)``, float64, which the + constrained channel sums to. + kernel + The interaction kernel, accepted by :func:`kernel_eligible`. + channel + The constrained channel, which is shifted uniformly over the charged + nodes of every frame to sum to its total charge, or -1 for none. + prefactor + Factor of every pair term, in eV Å. + communicator + The address of the MPI_Comm handle of a system decomposed into + domains, one per rank, as an int64 CPU tensor of one element, or zero, + the default, for a single domain. Across domains the call takes one + frame, whose charged nodes are those of every rank; the outputs are + those of the charged nodes of this rank, the other nodes carry no + charge, and ``virial`` is the share of this rank. + box + The cells of periodic frames with shape ``(F, 9)``, float64 on the + CPU, the rows of every row-major 3x3 cell being its lattice vectors, + or ``None`` for non-periodic frames; a zero row marks a non-periodic + frame. A periodic frame is evaluated on a mesh of spacing at most + ``spacing``, with B-splines of order ``order``. The device reads cells + in pinned memory in place while the stream runs the call, so the + caller keeps them alive and unchanged until then. A CUDA graph that + captures the call needs them pinned and reads them at every replay, + and it needs an eager call to have evaluated meshes of the same sizes + with the same channel count and ``order``. + order + The order of the B-splines of the mesh, at most :data:`MAX_ORDER`. + spacing + The largest mesh spacing, in Å. + + Returns + ------- + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, float64, half of the + interaction of the charges of the node with all others. + charge_gradient + Derivative of the energy with respect to ``charge`` with shape + ``(N, C)``, in eV per unit charge, through the constraint. + force + Force at fixed charges with shape ``(N, 3)``, in eV/Å. + atom_virial + Per-node virial at fixed charges with shape ``(N, 9)`` in eV: half of + every pair term to each of its nodes in a non-periodic frame; in a + periodic frame ``F (x) r`` with the rest of the virial of the cell + shared evenly over the charged nodes. + virial + Per-frame virial at fixed charges with shape ``(F, 9)`` in eV. + """ + ensure_registered() + code, coefficient = kernel_operator_arguments(kernel) + if communicator is None: + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") + return torch.ops.deepmd.latent_charge_sum( + coord, + charge, + node_mask, + n_node_per_frame, + total_charge, + box, + coefficient.to(coord.device), + communicator, + code, + int(channel), + float(prefactor), + int(order), + float(spacing), + ) diff --git a/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py b/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py index a11f5d49d6..d67b27eb6c 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py +++ b/deepmd/pt_expt/kernels/triton/sezm/grid_pair.py @@ -44,7 +44,11 @@ annotations, ) +from functools import ( + cache, +) from typing import ( + TYPE_CHECKING, Any, ) @@ -60,6 +64,13 @@ ) _Shape = tuple[int | torch.SymInt, ...] +_ShapeKey = tuple[str, int, int, int, int, bool, int, torch.dtype] +_LaunchConfig = tuple[int, int, int] + +if TYPE_CHECKING: + from triton.runtime.jit import ( + JITFunction, + ) __all__ = [ "GRID_PAIR_TRITON_AVAILABLE", @@ -879,8 +890,6 @@ def _unpack(value: Tensor, shape: _Shape, n_frames: int) -> Tensor: ) -_LAUNCH_CACHE: dict[tuple, tuple[int, int, int]] = {} - # Exact production-shape winners on RTX PRO 6000 Blackwell. The grid-pair # kernels hold one, two or three ``(P_PAD, C_BLK)`` accumulator tiles, so the # best channel width and grid tile change independently across differentiation @@ -1002,54 +1011,123 @@ def _unpack(value: Tensor, shape: _Shape, n_frames: int) -> Tensor: def _built_in_launch_config( - device_name: str, shape_key: tuple -) -> tuple[int, int, int] | None: + device_name: str, shape_key: _ShapeKey +) -> _LaunchConfig | None: """Return the built-in launch for one exact grid-pair shape.""" if device_name.startswith("NVIDIA RTX PRO 6000 Blackwell"): return _BLACKWELL_LAUNCH_CONFIGS.get(shape_key) return None +def _slot_partition(p_dim: int) -> tuple[int, int]: + """Split the coefficient axis into power-of-two contraction tiles.""" + # The slot axis is covered by the largest power of two below the count + # plus an optional low segment for the remainder, so 147 (degree six) + # pads to 128 + 32 and 75 (degree four) to 64 + 16 instead of the next + # power of two. The high segment keeps the tensor-core minimum of 16. + p_hi = max(16, 1 << (p_dim.bit_length() - 1)) + p_lo = max(16, _next_pow2(p_dim - p_hi)) if p_dim > p_hi else 0 + return p_hi, p_lo + + +def _launch_kwargs( + shape_key: _ShapeKey, config: _LaunchConfig +) -> dict[str, int | bool]: + """Share static geometry and tile parameters between validation and execution.""" + _, p_dim, c_per, n_focus, n_frames, packed, n_grid, _ = shape_key + p_hi, p_lo = _slot_partition(p_dim) + c_blk, block_g, stages = config + return { + "n_grid": n_grid, + "PACKED": packed, + "N_FOCUS": n_focus, + "N_FRAMES": n_frames, + "P_DIM": p_dim, + "P_HI": p_hi, + "P_LO": p_lo, + "C_ALL": c_per, + "C_BLK": c_blk, + "ALLOW_TF32": False, + "BLOCK_G": block_g, + "num_warps": 8 if c_blk >= 64 else 4, + "num_stages": stages, + } + + +@cache +def _select_launch_config( + kernel: JITFunction, + device: torch.device, + shape_key: _ShapeKey, + signature: tuple[torch.dtype | int, ...], + n_acc: int, +) -> _LaunchConfig | None: + """Compile and load candidates without allocating or executing their inputs. + + Tile resources depend on the static coefficient geometry and operand + dtypes. The leading pair count and physical strides only address rows; + their non-unit placeholders avoid specializing a dynamic batch size. + Triton's ``warmup`` accepts dtypes in place of tensor operands and + compiles the real device binary without dereferencing a pointer. Loading + its handles checks the device's resource limits before a configuration + enters the cache or a traced graph. Fake-tensor tracing therefore makes + the same choice as eager execution. + + The probe admits a tile configuration, not a byte-identical runtime binary. + Scalar and pointer-alignment specialization of the actual launch can still + produce a separate compiled artifact. + + The channel cap keeps the ``n_acc`` fp32 accumulator tiles register + resident. Exact shared-memory requirements are checked by Triton, whose + staging and transpose storage depend on the compiler and target device. + A swept configuration has priority but must pass the same validation. + ``None`` records that the reference composition must serve this layout. + """ + _, p_dim, c_per, *_ = shape_key + p_hi, p_lo = _slot_partition(p_dim) + p_eff = p_hi + p_lo + c_top = min(64, _next_pow2(c_per), max(16, _next_pow2(4096 // (n_acc * p_eff)))) + preferred = _built_in_launch_config(torch.cuda.get_device_name(device), shape_key) + candidates = [ + (c_blk, block_g, stages) + for c_blk in (c_top, 32, 16) + if c_blk <= c_top + for block_g, stages in ((64, 2), (32, 2), (32, 1), (16, 1)) + ] + if preferred is not None: + candidates.insert(0, preferred) + with torch.cuda.device(device): + for config in dict.fromkeys(candidates): + try: + compiled = kernel.warmup( + *signature, + 2, + grid=(1, 1), + **_launch_kwargs(shape_key, config), + ) + compiled._init_handles() + except triton.runtime.errors.OutOfResources: + continue + return config + return None + + def _launch( - kernel, + kernel: JITFunction, value: Tensor, n_grid: int, n_frames: int, packed: bool, - args: tuple, + args: tuple[Tensor | int | torch.SymInt, ...], n_acc: int, ) -> None: - """Launch with the largest tile the register and shared budgets admit. - - The channel block is capped so the ``n_acc`` fp32 accumulator tiles - ``(P_PAD, C_BLK)`` stay register-resident (a spilled accumulator is - read-modify-written through local memory once per grid block, which - dominated the three-accumulator second-order kernel before the cap). - The exact shared footprint additionally depends on Triton's internal - staging (dot operand buffers, transpose scratch), so candidates are - tried from the most to the least aggressive and the first that compiles - is cached per device and exact operator shape. Swept built-in launches - take precedence where available and fall back to the same compile search - if a later Triton version rejects one. - """ - n_batch = value.shape[0] - # The leading pair count follows the workload. The trailing coefficient - # layout and projector size define the Triton tile geometry and remain - # kernel-specialized. + """Record or execute a launch whose tile has passed device validation.""" n_grid = guard_int(n_grid) n_frames = guard_int(n_frames) coeff_dim, n_focus, packed_channels = (guard_int(size) for size in value.shape[1:]) - n_pair = n_batch * n_focus + n_pair = value.shape[0] * n_focus p_dim = coeff_dim * n_frames c_per = packed_channels // n_frames - # The slot axis is covered by the largest power of two below the count - # plus an optional low segment for the remainder, so 147 (degree six) - # pads to 128 + 32 and 75 (degree four) to 64 + 16 instead of the next - # power of two. The high segment keeps the tensor-core minimum of 16. - p_hi = max(16, 1 << (p_dim.bit_length() - 1)) - p_lo = max(16, _next_pow2(p_dim - p_hi)) if p_dim > p_hi else 0 - p_eff = p_hi + p_lo - c_top = min(64, _next_pow2(c_per), max(16, _next_pow2(4096 // (n_acc * p_eff)))) shape_key = ( kernel.fn.__name__, p_dim, @@ -1060,48 +1138,13 @@ def _launch( n_grid, value.dtype, ) - device_name = torch.cuda.get_device_name(value.device) - key = (device_name, *shape_key) - generic_candidates = [ - (c_blk, block_g, stages) - for c_blk in (c_top, 32, 16) - if c_blk <= c_top - for block_g, stages in ((64, 2), (32, 2), (32, 1), (16, 1)) - ] - if key in _LAUNCH_CACHE: - candidates = [_LAUNCH_CACHE[key]] - else: - built_in = _built_in_launch_config(device_name, shape_key) - candidates = ( - generic_candidates - if built_in is None - else [built_in, *(cfg for cfg in generic_candidates if cfg != built_in)] - ) - for c_blk, block_g, stages in candidates: - grid = (n_pair, (c_per + c_blk - 1) // c_blk) - try: - wrap_triton(kernel)[grid]( - *args, - n_pair, - n_grid=n_grid, - PACKED=packed, - N_FOCUS=n_focus, - N_FRAMES=n_frames, - P_DIM=p_dim, - P_HI=p_hi, - P_LO=p_lo, - C_ALL=c_per, - C_BLK=c_blk, - ALLOW_TF32=False, - BLOCK_G=block_g, - num_warps=8 if c_blk >= 64 else 4, - num_stages=stages, - ) - except triton.runtime.errors.OutOfResources: - continue - _LAUNCH_CACHE[key] = (c_blk, block_g, stages) - return - raise _NoViableConfig(p_dim, c_per) + signature = tuple(arg.dtype if isinstance(arg, Tensor) else 2 for arg in args) + config = _select_launch_config(kernel, value.device, shape_key, signature, n_acc) + if config is None: + raise _NoViableConfig(p_dim, c_per) + c_blk = config[0] + grid = (n_pair, (c_per + c_blk - 1) // c_blk) + wrap_triton(kernel)[grid](*args, n_pair, **_launch_kwargs(shape_key, config)) def _strides(*values: Tensor) -> _Shape: @@ -1156,9 +1199,9 @@ def _eager_packed( ) -> tuple[Tensor, ...]: """Eager einsum fallback on the packed layout. - Serves the shapes whose padded tiles exceed the shared-memory budget -- - the wide-slot fp32 regime of the exact-precision harnesses; the - production bf16 shapes never reach it. + Serves layouts for which no candidate tile fits the target device's + resources. The composition remains visible to the compiler when resource + validation rejects every fused configuration during tracing. """ tg = to_grid fg = from_grid diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_edge_cohort.py b/deepmd/pt_expt/kernels/triton/sezm/so2_edge_cohort.py new file mode 100644 index 0000000000..a6d9443bda --- /dev/null +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_edge_cohort.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Runtime edge-cohort backward for frozen, narrow SO(2) inference. + +Keep native GEMM column parallelism and arithmetic, but traverse all layers +of one cohort before advancing to the next cohort. Scratch scales with the +cohort, not the entire graph. The saved activations and returned input gradient +still have full edge extent. This Python launcher serves online compilation; +official CPU-traced exports retain the native deployable backward. +""" + +from __future__ import ( + annotations, +) + +import torch +import triton +from torch import ( + Tensor, +) +from torch.library import ( + wrap_triton, +) + +COHORT_EDGES = 8192 + + +def supports_edge_cohort( + grad_out: Tensor, z_all: Tensor, lmax: int, focus_dim: int, apply_alpha: bool +) -> bool: + """Restrict cohort scratch reuse to the validated FP32 Neo layout.""" + return ( + grad_out.is_cuda + and z_all.device == grad_out.device + and grad_out.dtype == torch.float32 + and z_all.dtype == torch.float32 + and grad_out.is_contiguous() + and z_all.is_contiguous() + and lmax == 3 + and focus_dim == 32 + and apply_alpha + and grad_out.ndim == 3 + and z_all.ndim == 4 + and grad_out.shape[1:] == (2, 320) + and z_all.shape[:2] == (2, 2) + and z_all.shape[2] == grad_out.shape[0] + and z_all.shape[3] == 320 + ) + + +def cohort_input_backward_impl( + grad_out: Tensor, + z_all: Tensor, + alpha: Tensor, + w0t: Tensor, + w1t: Tensor, + gw: Tensor, + gwt: Tensor, + cohort_edges: int, +) -> Tensor: + """Compute input gradients without retaining unrelated forward outputs. + + Kernels infer scratch focus strides from the cohort edge count. Independent + saved/output strides retain full-graph layout, including partial cohorts. + Only an internally owned residual is overwritten: matrix input ``gz`` and + GEMM output never alias, and caller-owned tensors remain read-only. + """ + from . import so2_value_path as sv + + if not supports_edge_cohort(grad_out, z_all, 3, 32, True): + raise ValueError("edge cohorts require contiguous FP32 Neo inference tensors") + if cohort_edges <= 0: + raise ValueError("cohort_edges must be positive") + if not all(t.is_contiguous() for t in (alpha, w0t, w1t, gw, gwt)): + raise ValueError("edge-cohort parameters must be contiguous") + n_gated, n_focus, n_edge, row = z_all.shape + result = grad_out.new_empty((n_focus, n_edge, row)) + if n_edge == 0: + return result + cap = min(n_edge, cohort_edges) + # A single cohort can reuse its owned output as the residual, just like + # native inference. Multiple cohorts need separate bounded scratch so + # completing a later cohort never overwrites an earlier output. + g = result if n_edge <= cap else grad_out.new_empty((n_focus, cap, row)) + gz = torch.empty_like(g) + sig = grad_out.new_empty((n_focus, cap, 96)) + bm, bn, bk, warps, stages = sv.stack_fp32_configs(32, 3)[2] + n_tiles = triton.cdiv(128, bn) + triton.cdiv(192, bn) + inside, point_config, recompute_config = sv._point_backward_schedule( + 32, 3, train=False + ) + for start in range(0, n_edge, cap): + count = min(cap, n_edge - start) + go = grad_out[start:] + a = alpha[start:] + z = z_all.view(-1)[start * row :] + grid = (triton.cdiv(count, bm) * n_tiles, n_focus) + sv._stack_gemm_bwd_kernel[grid]( + go, + go, + w0t, + w1t, + a, + g, + count, + n_gated, + L=3, + CF=32, + G_EDGE_MAJOR=True, + FOLD_ALPHA=True, + RES_IS_GZ=True, + BLOCK_M=bm, + BLOCK_N=bn, + BLOCK_K=bk, + num_warps=warps, + num_stages=stages, + ) + for layer in range(n_gated - 1, -1, -1): + if not inside: + rbm, rw, rs = recompute_config + sv._stack_recompute_kernel[(triton.cdiv(count, rbm), n_focus)]( + z, + gw, + sig, + count, + layer, + L=3, + CF=32, + BLOCK_M=rbm, + saved_edges=n_edge, + num_warps=rw, + num_stages=rs, + ) + pbm, pw, ps = point_config + sv._stack_point_bwd_kernel[(triton.cdiv(count, pbm), n_focus)]( + g, + z, + sig, + gwt, + gz, + sig, + g, + g, + count, + layer, + L=3, + CF=32, + GLOGIT_OUT=False, + GLOGIT_STORE=False, + RECOMPUTE_SIG=inside, + RECOVER_INPUT=False, + BLOCK_M=pbm, + saved_edges=n_edge, + num_warps=pw, + num_stages=ps, + ) + # Residual elements are loaded before their matching output store. + # Distinct (edge, focus, column) tiles own disjoint residual elements. + target = result.view(-1)[start * row :] if layer == 0 else g + sv._stack_gemm_bwd_kernel[grid]( + gz, + g, + w0t, + w1t, + gz, + target, + count, + layer, + L=3, + CF=32, + G_EDGE_MAJOR=False, + FOLD_ALPHA=False, + RES_IS_GZ=False, + BLOCK_M=bm, + BLOCK_N=bn, + BLOCK_K=bk, + output_edges=n_edge if layer == 0 else count, + num_warps=warps, + num_stages=stages, + ) + return result + + +@torch.library.custom_op("sezm_triton::so2_cohort_input_bwd", mutates_args=()) +def _cohort_input_backward( + grad_out: Tensor, + z_all: Tensor, + alpha: Tensor, + w0t: Tensor, + w1t: Tensor, + gw: Tensor, + gwt: Tensor, + cohort_edges: int, +) -> Tensor: + """Keep the runtime edge loop opaque without retaining x_local/u_final.""" + return cohort_input_backward_impl( + grad_out, z_all, alpha, w0t, w1t, gw, gwt, cohort_edges + ) + + +@_cohort_input_backward.register_fake +def _cohort_input_backward_fake( + grad_out: Tensor, + z_all: Tensor, + alpha: Tensor, + w0t: Tensor, + w1t: Tensor, + gw: Tensor, + gwt: Tensor, + cohort_edges: int, +) -> Tensor: + """Preserve full input-gradient shape under FakeTensor tracing.""" + return grad_out.new_empty(z_all.shape[1:]) + + +@torch.library.triton_op("sezm_triton::so2_cohort_alpha_bwd", mutates_args=()) +def _cohort_alpha_backward(grad_out: Tensor, x_local: Tensor, alpha: Tensor) -> Tensor: + """Expose alpha's short-lived output dependency to the graph scheduler.""" + from . import so2_value_path as sv + + n_edge, n_focus = alpha.shape + result = torch.empty_like(alpha) + if sv._has_no_edges(n_edge): + return result + bm, warps, stages = sv.gate_config(32, 3) + wrap_triton(sv._stack_grad_alpha_kernel)[(triton.cdiv(n_edge, bm), n_focus)]( + grad_out, + x_local, + alpha, + result, + n_edge, + L=3, + CF=32, + BLOCK_M=bm, + num_warps=warps, + num_stages=stages, + ) + return result + + +def cohort_backward( + grad_out: Tensor, + x_local: Tensor, + z_all: Tensor, + u_final: Tensor, + alpha: Tensor, + w0t: Tensor, + w1t: Tensor, + gw: Tensor, + gwt: Tensor, + lmax: int, + focus_dim: int, + apply_alpha: bool, +) -> tuple[Tensor, Tensor]: + """First-order dispatcher; training/higher derivatives never enter here.""" + del u_final + if not supports_edge_cohort(grad_out, z_all, lmax, focus_dim, apply_alpha): + raise ValueError("unsupported edge-cohort layout") + grad_alpha = _cohort_alpha_backward(grad_out, x_local, alpha) + # make_fx/AOT may trace with a small example and retain that route for + # larger inputs. Keep edge partitioning inside the opaque runtime operator + # so its loop always uses the actual edge count, not the tracing example. + grad_u0 = _cohort_input_backward( + grad_out, z_all, alpha, w0t, w1t, gw, gwt, COHORT_EDGES + ) + return grad_u0, grad_alpha diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py index 7a8ac3fd43..5f2ab9f725 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py @@ -1369,6 +1369,7 @@ def _stack_recompute_kernel( L: tl.constexpr, CF: tl.constexpr, BLOCK_M: tl.constexpr, + saved_edges=0, ): """Recompute the gate sigmoids from the saved pre-activation (backward).""" ROW: tl.constexpr = (3 * L + 1) * CF @@ -1386,7 +1387,8 @@ def _stack_recompute_kernel( cm = mm & (nc < CF)[None, :] wm = ((nc < CF)[:, None]) & ((nc < CF)[None, :]) - z_row = z_ptr + (layer * n_focus + fid) * n_edge * ROW + offs_m * ROW + z_stride = tl.where(saved_edges > 0, saved_edges, n_edge).to(tl.int64) + z_row = z_ptr + (layer * n_focus + fid) * z_stride * ROW + offs_m * ROW z_s = tl.load(z_row[:, None] + nc[None, :], mask=cm, other=0.0).to(tl.float32) sig_row = sig_ptr + (fid * n_edge + offs_m) * LG for g in tl.static_range(L): @@ -1420,6 +1422,7 @@ def _stack_point_bwd_kernel( RECOMPUTE_SIG: tl.constexpr, RECOVER_INPUT: tl.constexpr, BLOCK_M: tl.constexpr, + saved_edges=0, ): """Pointwise part of the gated-layer backward. @@ -1456,7 +1459,8 @@ def _stack_point_bwd_kernel( wm = ((nc < CF)[:, None]) & ((nc < CF)[None, :]) g_row = g_ptr + fid * n_edge * ROW + offs_m * ROW - z_row = z_ptr + (layer * n_focus + fid) * n_edge * ROW + offs_m * ROW + z_stride = tl.where(saved_edges > 0, saved_edges, n_edge).to(tl.int64) + z_row = z_ptr + (layer * n_focus + fid) * z_stride * ROW + offs_m * ROW gz_row = gz_ptr + fid * n_edge * ROW + offs_m * ROW sig_row = sig_ptr + (fid * n_edge + offs_m) * LG gl_row = gl_ptr + (fid * n_edge + offs_m) * LG @@ -2014,6 +2018,7 @@ def _stack_gemm_bwd_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + output_edges=0, ): """Backward block GEMM ``g_u = residual + gz @ W^T`` over both blocks. @@ -2044,7 +2049,8 @@ def _stack_gemm_bwd_kernel( gz_row = gz_ptr + fid * ROW + offs_m * (n_focus * ROW) else: gz_row = gz_ptr + fid * n_edge * ROW + offs_m * ROW - gu_row = gu_ptr + fid * n_edge * ROW + offs_m * ROW + output_stride = tl.where(output_edges > 0, output_edges, n_edge).to(tl.int64) + gu_row = gu_ptr + fid * output_stride * ROW + offs_m * ROW if FOLD_ALPHA: alpha = tl.load(alpha_ptr + offs_m * n_focus + fid, mask=m_mask, other=0.0) @@ -2435,6 +2441,7 @@ def _launch_stack_point_backward( L=lmax, CF=focus_dim, BLOCK_M=block_m, + saved_edges=0, num_warps=warps, num_stages=stages, ) @@ -2458,6 +2465,7 @@ def _launch_stack_point_backward( RECOMPUTE_SIG=recompute_sigmoid, RECOVER_INPUT=recover, BLOCK_M=block_m, + saved_edges=0, num_warps=warps, num_stages=stages, ) @@ -3191,14 +3199,18 @@ def _stack_backward_traversal( if keep and n_gated > 0 else torch.empty((n_focus, n_edge, row), device=device, dtype=dtype) ) - # Inference retains no per-layer upstream state. Two buffers carry the - # reverse recurrence without exposing an unrolled layer stack for - # functionalization into select-scatter copies. - g_spare = ( - torch.empty((n_focus, n_edge, row), device=device, dtype=dtype) - if not keep and n_gated > 0 - else None - ) + # Without weight gradients or saved per-layer state, the owned residual + # buffer can also hold the next layer's gradient. The GEMM reads its matrix + # input from the separate gz buffer; each output tile reads only its own + # residual elements before overwriting them. Never alias gz or caller inputs. + # Weight-gradient and second-order traversals retain their original buffers. + g_spare = None + if not keep and n_gated > 0: + g_spare = ( + torch.empty((n_focus, n_edge, row), device=device, dtype=dtype) + if with_weights + else g_cur + ) wrap_triton(_stack_gemm_bwd_kernel)[ (triton.cdiv(n_edge, block_m) * n_tiles, n_focus) ]( @@ -3218,6 +3230,7 @@ def _stack_backward_traversal( BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, + output_edges=0, num_warps=warps, num_stages=stages, ) @@ -3402,6 +3415,7 @@ def _stack_backward_traversal( BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, + output_edges=0, num_warps=warps, num_stages=stages, ) @@ -4956,7 +4970,25 @@ def _mixing_stack_backward(ctx, grad_out, grad_z, grad_u): None, None, ) - grad_u0, grad_alpha = _mixing_stack_bwd_op( + backward_op = _mixing_stack_bwd_op + # The fused FP32 stack is selected by the existing Triton inference + # controls. Only its validated layout uses cohort-sized scratch; weight, + # saved-state and higher-order requests have already taken the train path. + from .so2_edge_cohort import ( + cohort_backward, + supports_edge_cohort, + ) + + if supports_edge_cohort( + grad_out, z_all, ctx.lmax, ctx.focus_dim, ctx.apply_alpha + ) and all( + value.dtype == torch.float32 + and value.device == grad_out.device + and value.is_contiguous() + for value in (alpha, w0_all, w1_all, gw_all) + ): + backward_op = cohort_backward + grad_u0, grad_alpha = backward_op( grad_out, x_local, z_all, diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py index c6e54bd17f..4369b6ddfd 100644 --- a/deepmd/pt_expt/kernels/utils.py +++ b/deepmd/pt_expt/kernels/utils.py @@ -180,7 +180,7 @@ def cuda_infer_level() -> int: - DPA1 (``se_atten``): the descriptor mega kernels (:mod:`.cuda.dpa1.graph_descriptor`) serve the concat, attention-free graph lower, and the energy fitting runs through the - fused cuBLAS network (:mod:`.cuda.graph_fitting`). + fused fitting network (:mod:`.graph_fitting`). - DPA4 (``sezm``): the operators whose profit is memory traffic and therefore holds on every part measured -- the fused SO(3) grid pair product (:mod:`.cuda.dpa4.grid_pair`), which keeps the grid field in @@ -190,7 +190,7 @@ def cuda_infer_level() -> int: (:mod:`.cuda.dpa4.edge_radial`), which evaluates that chain once instead of once per consumer. - All graph-lowered models: the force / virial assembly scatters - through :mod:`.cuda.edge_force_virial`. + through :mod:`.edge_force_virial`. - ``2`` -- adds the end-to-end energy-force operator: a graph-lowered energy model whose descriptor and fitting are both fused-eligible collapses its diff --git a/deepmd/pt_expt/long_range/__init__.py b/deepmd/pt_expt/long_range/__init__.py new file mode 100644 index 0000000000..9c230575e7 --- /dev/null +++ b/deepmd/pt_expt/long_range/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + LESKernel, + SOGKernel, +) +from .latent_charge import ( + LatentChargeLongRange, +) + +__all__ = [ + "BaseLongRange", + "LESKernel", + "LatentChargeLongRange", + "SOGKernel", +] diff --git a/deepmd/pt_expt/long_range/base_long_range.py b/deepmd/pt_expt/long_range/base_long_range.py new file mode 100644 index 0000000000..66135372f9 --- /dev/null +++ b/deepmd/pt_expt/long_range/base_long_range.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import torch + +from deepmd.dpmodel.long_range import ( + make_base_long_range, +) + +BaseLongRange = make_base_long_range(torch.Tensor) diff --git a/deepmd/pt_expt/long_range/kernel.py b/deepmd/pt_expt/long_range/kernel.py new file mode 100644 index 0000000000..3675e6b988 --- /dev/null +++ b/deepmd/pt_expt/long_range/kernel.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.long_range.kernel import LESKernel as LESKernelDP +from deepmd.dpmodel.long_range.kernel import SOGKernel as SOGKernelDP +from deepmd.pt_expt.common import ( + register_dpmodel_mapping, + to_torch_array, +) + + +def _torch_kernel(kernel_dp: type) -> type: + """Wrap a dpmodel interaction kernel as a module whose parameters are trained. + + The attributes listed in ``PARAMETER_NAMES`` are stored as + ``torch.nn.Parameter`` that require a gradient exactly when the kernel is + trainable, both at construction and when a deserialized value is assigned. + """ + + class TorchKernel(kernel_dp, torch.nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + torch.nn.Module.__init__(self) + kernel_dp.__init__(self, *args, **kwargs) + + def __call__(self, *args: Any, **kwargs: Any) -> torch.Tensor: + return torch.nn.Module.__call__(self, *args, **kwargs) + + def forward(self, r_sq: torch.Tensor) -> torch.Tensor: + return self.call(r_sq) + + def __setattr__(self, name: str, value: Any) -> None: + if name in self.PARAMETER_NAMES and "_parameters" in self.__dict__: + if not isinstance(value, torch.nn.Parameter): + value = torch.nn.Parameter( + to_torch_array(value).detach().clone(), + requires_grad=self.trainable, + ) + if name in self._parameters: + self._parameters[name] = value + return + super().__setattr__(name, value) + + TorchKernel.__name__ = TorchKernel.__qualname__ = kernel_dp.__name__ + register_dpmodel_mapping( + kernel_dp, lambda v, cls=TorchKernel: cls.deserialize(v.serialize()) + ) + return TorchKernel + + +LESKernel = _torch_kernel(LESKernelDP) +SOGKernel = _torch_kernel(SOGKernelDP) diff --git a/deepmd/pt_expt/long_range/latent_charge.py b/deepmd/pt_expt/long_range/latent_charge.py new file mode 100644 index 0000000000..cadd860d50 --- /dev/null +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import torch + +from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, +) +from deepmd.dpmodel.long_range.latent_charge import ( + LatentChargeLongRange as LatentChargeLongRangeDP, +) +from deepmd.pt_expt.common import ( + torch_module, +) +from deepmd.pt_expt.fitting.latent_charge import ( + LatentChargeFitting, +) + +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + LESKernel, + SOGKernel, +) + + +@BaseLongRange.register("latent_charge") +@torch_module +class LatentChargeLongRange(LatentChargeLongRangeDP): + """Energy of interacting latent charges, periodic or not (pt_expt).""" + + # The dpmodel class builds the charge network and the kernel, which the + # assignment converts to these pt_expt modules; importing their classes + # registers the conversions. + charge_net: LatentChargeFitting + kernel: LESKernel | SOGKernel + + def _pair_kernel( + self, r_sq: torch.Tensor, src: torch.Tensor, n: int + ) -> torch.Tensor: + # A graph traced for a symbolic pair count plans the reduction of the + # coefficient gradients over the pairs before it sees a batch. In + # training the coefficients therefore reach the pairs through the node + # axis, which turns that reduction into a scatter over the pairs and a + # sum over the nodes, parallel in the pair count of every batch. A + # concrete pair count is reduced directly. + if not (self.training and isinstance(r_sq.shape[0], torch.SymInt)): + return self.kernel(r_sq) + coefficients = { + name: value.expand(n, *value.shape)[src] + for name, value in self.kernel.coefficients(r_sq).items() + } + return self.kernel.evaluate(r_sq, coefficients) + + def operator_eligible(self) -> bool: + """Whether :meth:`energy_force_from_coord` evaluates this module. + + The operator evaluates the kernels, channel counts and B-spline orders + it is instantiated for. The compact lower that calls it reads no frame + input: the charges depend on the descriptor alone, and a constrained + channel sums to zero or to the total charge of a folded charge state. + """ + from deepmd.pt_expt.kernels.long_range import ( + MAX_CHANNELS, + MAX_ORDER, + kernel_eligible, + op_available, + ) + + return ( + op_available() + and kernel_eligible(self.kernel) + and self.charge_net.dim_out <= MAX_CHANNELS + and self.pme_order <= MAX_ORDER + and self.charge_net.numb_fparam == 0 + and self.charge_net.numb_aparam == 0 + and not self.uses_charge_spin() + ) + + def operator_constraint( + self, n_frames: int, like: torch.Tensor + ) -> tuple[torch.Tensor, int]: + """The total charge of every frame and the constrained channel of the operator. + + Parameters + ---------- + n_frames + The number of frames. + like + A tensor giving the device of the total charges. + + Returns + ------- + tuple[torch.Tensor, int] + The total charges with shape ``(n_frames,)``, float64 and + contiguous, and the constrained channel, or -1 without a + constraint, in which case the total charges are zero. + """ + total_charge = like.new_zeros(n_frames, dtype=torch.float64) + if self.charge_constraint is None: + return total_charge, -1 + total_charge = self._total_charge(n_frames, total_charge, None, None) + return total_charge.contiguous(), int(self.charge_constraint["channel"]) + + def energy_force_from_coord( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + node_mask: torch.Tensor, + n_node: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the module for inference from the positions of the nodes. + + The operator sums a non-periodic frame directly over its pairs, which + it enumerates itself, so that the memory is linear in the node count, + and a periodic frame on the mesh of its cell. It applies the + constraint and returns the force at fixed charges and the derivative + of the energy with respect to the charges the network predicts; + pulled back through the charge network, that derivative is the + derivative of the energy with respect to the descriptor, through which + the caller adds the response of the charges to the positions. + + Parameters + ---------- + descriptor + Invariant node features with shape ``(N, dim_descrpt)``. + atype + Node types with shape ``(N,)``. + node_mask + Nodes carrying charges with shape ``(N,)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + coord + Node positions with shape ``(N, 3)`` in Å. + communicator + The MPI communicator of a frame decomposed into domains, one per + rank, whose charged nodes are those of every rank, in the form of + :func:`~deepmd.pt_expt.kernels.long_range.latent_charge_sum`; + ``None`` for a single domain. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the + CPU, a zero row marking a non-periodic frame, or ``None`` when no + frame is periodic. A module that does not serve periodic systems + evaluates every frame as non-periodic. + + Returns + ------- + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, float64. + descriptor_gradient + Derivative of the energy with respect to the descriptor through + the charges, with the shape and precision of ``descriptor``. + force + Force at fixed charges with shape ``(N, 3)`` in eV/Å, float32. + atom_virial + Per-node virial at fixed charges with shape ``(N, 9)`` in eV, + float32, each pair term split half to each of its nodes. + virial + Per-frame virial at fixed charges with shape ``(nf, 9)`` in eV, + float32; across domains, the share of the nodes of this rank. + """ + from deepmd.pt_expt.kernels.long_range import ( + latent_charge_sum, + ) + + leaf = descriptor.detach().requires_grad_(True) + with torch.enable_grad(): + charge = self.charge_net.call_graph(leaf, atype)["latent_charge"] + total_charge, channel = self.operator_constraint(n_node.shape[0], coord) + atom_energy, charge_gradient, force, atom_virial, virial = latent_charge_sum( + coord.to(torch.float64).contiguous(), + charge.detach().to(torch.float32).contiguous(), + node_mask.contiguous(), + n_node, + total_charge, + self.kernel, + channel, + COULOMB_CONSTANT, + communicator, + # A module that does not serve periodic systems sums every frame + # directly. + box if self.periodic else None, + self.pme_order, + self.mesh_spacing() if self.periodic else 1.0, + ) + (descriptor_gradient,) = torch.autograd.grad( + charge, leaf, charge_gradient.to(charge.dtype) + ) + return ( + atom_energy, + descriptor_gradient.to(descriptor.dtype), + force, + atom_virial, + virial, + ) diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index a4a497d453..f149053489 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -1,10 +1,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Autograd assembly: graph energy -> force/virial/atom_virial via grad(E, edge_vec). -torch-only. The pure-array scatter (edge_force_virial) is shared with dpmodel; -this module supplies the single backward pass that produces g_e = dE/d(edge_vec). +torch-only. The pure-array scatters (edge_force_virial, pair_force_virial) are +shared with dpmodel; this module supplies the single backward pass that +produces the gradient of the energy with respect to every geometry leaf: the +edge vectors of the neighbor graph and the geometry leaves of the topologies. """ +import dataclasses +from collections.abc import ( + Sequence, +) + import torch from deepmd.dpmodel import ( @@ -19,9 +26,14 @@ get_deriv_name_mag, ) from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, + MeshGrid, NeighborGraph, + PairGraph, edge_force_virial, frame_id_from_n_node, + node_validity_mask, + pair_force_virial, segment_sum, ) from deepmd.pt.utils import ( @@ -44,6 +56,91 @@ ) +def releaf_topologies( + topologies: dict[str, PairGraph | KGrid | MeshGrid], +) -> dict[str, PairGraph | KGrid | MeshGrid]: + """Make the geometry fields of every topology autograd leaves. + + The leaves are made from the inputs of a traced graph lower, so that the + backward roots at its inputs: a leaf created inside a ``make_fx`` trace + has a NaN gradient under inductor. + """ + return { + name: dataclasses.replace( + topology, + **{ + leaf: getattr(topology, leaf).detach().requires_grad_(True) + for leaf in topology.LEAVES + }, + ) + for name, topology in topologies.items() + } + + +def _cell_force_virial( + g_coord: torch.Tensor, + g_strain: torch.Tensor, + coord: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor | None, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + """Force and virial of a reciprocal-space energy from its geometry leaves. + + The gradient with respect to the node positions is the negative force, + and that with respect to the zero strain of every frame the negative + virial. The atomic virial takes ``F_i (x) r_i`` at every node, the force + index first as for the edge virial, and shares the rest of the virial of + the frame, the part of the cell, evenly over its owned nodes, so that the + nodes of a frame sum to its virial; the halo nodes, images that carry no + charge, take no share. + + Parameters + ---------- + g_coord + (N, 3) gradient of the energy with respect to the positions. + g_strain + (nf, 3, 3) gradient of the energy with respect to the strain. + coord + (N, 3) node positions. + n_node + (nf,) per-frame node counts. + n_local + (nf,) per-frame owned node counts, or ``None`` when every node is + owned. + do_atomic_virial + Whether to return the per-node virial. + + Returns + ------- + force + (N, 3) per-node force. + atom_virial + (N, 3, 3) per-node virial, or ``None``. + virial + (nf, 3, 3) per-frame virial. + """ + force = -g_coord + virial = -g_strain + if not do_atomic_virial: + return force, None, virial + n_total = force.shape[0] + frame_id = frame_id_from_n_node(n_node, n_total) + atom_virial = force[:, :, None] * coord.to(force.dtype)[:, None, :] + remainder = virial - segment_sum(atom_virial, frame_id, n_node.shape[0]) + if n_local is None: + owned = node_validity_mask(n_node, n_total) + count = n_node + else: + owned = node_ownership_mask(n_node, n_local, n_total) + count = n_local + share = owned.to(force.dtype) / torch.clamp( + count.to(force.dtype)[frame_id], min=1.0 + ) + atom_virial = atom_virial + remainder[frame_id] * share[:, None, None] + return force, atom_virial, virial + + def edge_energy_deriv( energy: torch.Tensor, edge_vec: torch.Tensor, @@ -60,6 +157,8 @@ def edge_energy_deriv( do_atomic_virial: bool = False, create_graph: bool = False, force_precision: torch.dtype | None = None, + topologies: Sequence[PairGraph | KGrid | MeshGrid] = (), + n_local: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """Return (force, atom_virial_or_None, virial) from a graph energy. @@ -79,6 +178,15 @@ def edge_energy_deriv( reduction. A retained graph (training / double backward) keeps the fp64 leaf dtype so training numerics are unchanged. + The same backward differentiates the energy with respect to the geometry + leaves of every topology in ``topologies``, and their pullbacks are summed + in: the energy reaches the coordinates through all its geometry leaves, so + its complete force and virial are the sum of their pullbacks. The pair + vectors of a pair topology are pulled back by the shared + pair_force_virial scatter; a reciprocal-space topology carries the node + positions, whose gradient is its force, and a zero strain, whose gradient + is its negative virial (see :func:`_cell_force_virial`). + Parameters ---------- energy @@ -109,6 +217,12 @@ def edge_energy_deriv( whether to materialize the per-atom virial (else ``None`` is returned). create_graph whether the backward retains a graph (training, for second-order grad). + topologies + Topologies on the same node axis whose ``LEAVES`` are further + autograd leaves of ``energy``. + n_local + (nf,) per-frame owned node counts, or ``None`` when every node is + owned; the owned nodes share the virial of a cell. Returns ------- @@ -119,9 +233,12 @@ def edge_energy_deriv( virial (nf, 3, 3) per-frame virial (always computed). """ - (g_e,) = torch.autograd.grad( + leaves = [ + getattr(topology, name) for topology in topologies for name in topology.LEAVES + ] + g_e, *g_leaves = torch.autograd.grad( energy.sum() if energy.dim() else energy, - edge_vec, + (edge_vec, *leaves), create_graph=create_graph, retain_graph=True, ) @@ -197,6 +314,29 @@ def edge_energy_deriv( force, atom_virial, virial = edge_force_virial( g_e, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity ) + offset = 0 + for topology in topologies: + grads = g_leaves[offset : offset + len(topology.LEAVES)] + offset += len(topology.LEAVES) + if isinstance(topology, PairGraph): + (g_p,) = grads + leaf_force, leaf_atom_virial, leaf_virial = pair_force_virial( + g_p, + topology.vec, + topology.index, + topology.mask, + n_node, + node_capacity=force.shape[0], + ) + else: + g_coord, g_strain = grads + leaf_force, leaf_atom_virial, leaf_virial = _cell_force_virial( + g_coord, g_strain, topology.coord, n_node, n_local, do_atomic_virial + ) + force = force + leaf_force.to(force.dtype) + virial = virial + leaf_virial.to(virial.dtype) + if do_atomic_virial: + atom_virial = atom_virial + leaf_atom_virial.to(atom_virial.dtype) return force, (atom_virial if do_atomic_virial else None), virial @@ -211,6 +351,7 @@ def fit_output_to_model_output_graph( n_local: torch.Tensor | None = None, force_precision: torch.dtype | None = None, spin_leaf: torch.Tensor | None = None, + topologies: dict[str, PairGraph | KGrid | MeshGrid] | None = None, ) -> dict[str, torch.Tensor]: """Graph analogue of the dense pt_expt ``fit_output_to_model_output``. @@ -282,6 +423,10 @@ def fit_output_to_model_output_graph( ``grad(E, [edge_vec, spin])`` jointly in one call; here the two grads are separate so :func:`edge_energy_deriv`'s signature stays untouched. ``None`` (default): unchanged, no mag output. + topologies + Topologies read by the atomic model, by name. Their geometry leaves + are autograd leaves of ``fit_ret`` next to ``graph.edge_vec``; force + and virial include their pullbacks (see :func:`edge_energy_deriv`). Returns ------- @@ -393,6 +538,8 @@ def reduce_by_frame(data: torch.Tensor) -> torch.Tensor: do_atomic_virial=(vdef.c_differentiable and do_atomic_virial), create_graph=create_graph, force_precision=force_precision if not create_graph else None, + topologies=() if topologies is None else tuple(topologies.values()), + n_local=n_local, ) # force (N, 3) -> (N, 1, 3) [flat; caller unravels at I/O boundary] ff_list.append(force.reshape(N, 1, 3)) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index f7e021a684..f18d3e33fc 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy +import dataclasses from typing import ( Any, ) @@ -18,12 +19,23 @@ from deepmd.dpmodel.model.make_hessian_model import ( make_hessian_model, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + flatten_topologies, +) from deepmd.dpmodel.utils.neighbor_list import ( NeighborList, ) +# Importing the module also registers the freeze-time measurement of the +# fitting-route tables, which runs before an export traces the compact +# canonical lower. +from deepmd.pt_expt.kernels.fitting_routes import ( + fitting_routes, +) + from .make_model import ( _translate_energy_keys, + check_owned_frames, fused_atom_bias, make_model, ) @@ -138,6 +150,9 @@ def forward_lower_canonical_graph( *, do_atomic_virial: bool, spin: torch.Tensor | None = None, + coord: torch.Tensor | None = None, + communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Evaluate an eligible compressed canonical deployment graph. @@ -150,21 +165,41 @@ def forward_lower_canonical_graph( n_local Per-frame owned node counts with shape ``(nf,)``, int64. source - Source-node indices with shape ``(S,)``, uint32. + Source-node indices with shape ``(S,)``, uint32 or int64, or their + exported byte carrier with shape ``(S, W)`` for ``W`` equal to + four or eight bytes. edge_vec Destination-major edge vectors with shape ``(S, 3)``, float32. destination_row_ptr Destination CSR offsets with shape ``(N + 1,)``, int64. source_row_ptr - Source CSR offsets with shape ``(N + 1,)``, int64. + Source CSR offsets with shape ``(N + 1,)``, int64. A DPA4C model + evaluated by a fused operator, which every DPA4C model without a + long-range module is, also accepts a two-entry placeholder, which + declines the per-atom virial; that output is then zero. source_order - Source-grouped edge positions with shape ``(S,)`` and the same - dtype as ``source``. + Source-grouped edge positions with shape ``(S,)``, int64, or a + two-entry placeholder alongside a placeholder ``source_row_ptr``. do_atomic_virial Whether to return the per-node virial. spin Per-node magnetic moments with shape ``(N, 3)`` for a native-spin model, or ``None``. Ghost rows carry their owner's moment. + coord + Node positions with shape ``(N, 3)`` in Å, float64, for a model + with a long-range module; ``None`` otherwise. + communicator + For a model with a long-range module, the address of the MPI_Comm + handle of a frame decomposed into domains, one per rank, as an + int64 CPU tensor of one element, whose owned nodes are then those + of one domain; zero or ``None`` for a complete frame. + box + For a model with a long-range module, the cells of the frames + with shape ``(nf, 9)``, float64 on the CPU, the rows of every + row-major 3x3 cell its lattice vectors; a zero row marks a + complete non-periodic frame, and ``None`` non-periodic frames. + The halo nodes of a periodic frame are images of its owned + nodes. Returns ------- @@ -201,32 +236,106 @@ def forward_lower_canonical_graph( source_order=source_order, ) validate_canonical_graph_shapes(graph, atype.shape[0]) - atype, output_mask = self.atomic_model._prepare_graph_nodes( - n_node, - n_local, - atype, - edge_vec, - ) learned, pair_potential = self.atomic_model.fused_decomposition() descriptor = learned.descriptor fitting = learned.fitting_net atom_bias = fused_atom_bias(self.atomic_model, learned) - if use_dpa4c: - energy, atom_energy, force, virial, atom_virial, force_mag = ( + + def fused_graph() -> CanonicalGraph: + # A program without the per-atom virial declines it with the + # two-entry placeholders of the source CSR, so that a fused + # operator stores no edge cotangent whatever its caller passes. + if do_atomic_virial: + return graph + return dataclasses.replace( + graph, + source_row_ptr=source_row_ptr.new_zeros(2), + source_order=source_order.new_zeros(2), + ) + + if learned.has_long_range(): + from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + dpa4c_canonical_compress_long_range_energy_force, + dpa4c_canonical_long_range_energy_force, + fused_long_range_eligible, + ) + + if coord is None: + raise ValueError( + "a model with a long-range module reads the node positions" + ) + check_owned_frames(n_node, n_local, communicator, box) + force_mag = None + if fused_long_range_eligible(learned.long_range): + # The fused operator derives the ownership mask from the frame + # counts itself, and the fitting-route tables measured for the + # energy fitting and the charge network become literals of the + # traced program, as for the operator without a long-range + # module. + energy, atom_energy, force, virial, atom_virial, mask = ( + dpa4c_canonical_compress_long_range_energy_force( + descriptor, + fitting, + learned.long_range, + fused_graph(), + atype, + coord, + communicator, + box, + atom_bias, + fitting_routes(fitting), + fitting_routes(learned.long_range.charge_net), + ) + ) + else: + atype, output_mask = self.atomic_model._prepare_graph_nodes( + n_node, + n_local, + atype, + edge_vec, + ) + mask = output_mask.to(torch.int32) + energy, atom_energy, force, virial, atom_virial = ( + dpa4c_canonical_long_range_energy_force( + descriptor, + fitting, + learned.long_range, + graph, + atype, + coord, + communicator, + box, + output_mask, + atom_bias, + do_atomic_virial, + ) + ) + elif use_dpa4c: + # The fused operator derives the ownership mask from the frame + # counts itself, so nothing runs on the node axis around it. The + # fitting-route table measured for the network on the backend + # device, if any, becomes a literal of the traced program. + energy, atom_energy, force, virial, atom_virial, force_mag, mask = ( dpa4c_canonical_compress_energy_force( descriptor, fitting, - graph, + fused_graph(), atype, - output_mask, atom_bias, - do_atomic_virial, spin, None if pair_potential is None else pair_potential.pair_table, + fitting_routes(fitting), ) ) else: + atype, output_mask = self.atomic_model._prepare_graph_nodes( + n_node, + n_local, + atype, + edge_vec, + ) force_mag = None + mask = output_mask.to(torch.int32) energy, atom_energy, force, virial, atom_virial = ( dpa1_canonical_compress_energy_force( descriptor, @@ -246,7 +355,7 @@ def forward_lower_canonical_graph( "energy": energy, "force": force, "virial": virial, - "mask": output_mask.to(torch.int32), + "mask": mask, } if spin is not None: if force_mag is None or force_mag.ndim != 2: @@ -271,6 +380,9 @@ def forward_lower_canonical_graph_exportable( destination_row_ptr: torch.Tensor, source_row_ptr: torch.Tensor, source_order: torch.Tensor, + coord: torch.Tensor | None = None, + communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, *, do_atomic_virial: bool, **make_fx_kwargs: Any, @@ -283,6 +395,12 @@ def forward_lower_canonical_graph_exportable( Compact graph node and edge tensors. destination_row_ptr, source_row_ptr, source_order Compact dual-CSR topology. + coord, communicator, box + For a model with a long-range module, sample node positions with + shape ``(N, 3)``, float64, a sample communicator, an int64 CPU + tensor of one element, and sample cells with shape ``(nf, 9)``, + float64 on the CPU, which the traced module then takes as its + ninth to eleventh inputs; ``None`` otherwise. do_atomic_virial Whether the traced output includes per-node virial. **make_fx_kwargs @@ -291,9 +409,59 @@ def forward_lower_canonical_graph_exportable( Returns ------- torch.nn.Module - Traced eight-input compact deployment module. + Traced compact deployment module of eight inputs, or eleven with + the node positions, the communicator and the cells. """ model = self + if not (coord is None) == (communicator is None) == (box is None): + raise ValueError( + "the compact lower of a model with a long-range module takes " + "the node positions, the communicator and the cells together" + ) + + if coord is not None: + + def fn_long_range( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + source: torch.Tensor, + edge_vec: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_row_ptr: torch.Tensor, + source_order: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor, + box: torch.Tensor, + ) -> dict[str, torch.Tensor]: + return model.forward_lower_canonical_graph( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + do_atomic_virial=do_atomic_virial, + coord=coord, + communicator=communicator, + box=box, + ) + + return make_fx(fn_long_range, **make_fx_kwargs)( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + coord, + communicator, + box, + ) def fn( atype: torch.Tensor, @@ -613,6 +781,7 @@ def forward_lower_graph_exportable( do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, destination_sorted: bool = False, + topologies: dict[str, Any] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_common_lower_graph`` into an exportable module with @@ -646,6 +815,11 @@ def forward_lower_graph_exportable( and ``destination_order`` is identity. fparam, aparam, do_atomic_virial, charge_spin As in ``forward_lower``. + topologies + The sample topologies the atomic model declares, or ``None``; + their fields follow ``charge_spin`` as further inputs, in the + layout of + :func:`~deepmd.dpmodel.utils.neighbor_graph.flatten_topologies`. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -680,6 +854,7 @@ def forward_lower_graph_exportable( charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, destination_sorted=destination_sorted, + topologies=topologies, **make_fx_kwargs, ) @@ -688,6 +863,67 @@ def forward_lower_graph_exportable( do_grad_r = self.do_grad_r("energy") do_grad_c = self.do_grad_c("energy") + if topologies: + # The topology inputs follow ``charge_spin``, their fields in the + # layout of flatten_topologies. + + def fn_topologies( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + *topology_vals: torch.Tensor, + ) -> dict[str, torch.Tensor]: + model_ret = traced( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + *topology_vals, + ) + return _translate_energy_keys( + model_ret, + do_grad_r=do_grad_r, + do_grad_c=do_grad_c, + do_atomic_virial=do_atomic_virial, + local=True, + ) + + return make_fx(fn_topologies, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + *flatten_topologies(topologies), + ) + def fn( atype: torch.Tensor, n_node: torch.Tensor, diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 2071875b5c..63e1321bdc 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -31,6 +31,9 @@ from deepmd.pt_expt.fitting import ( BaseFitting, ) +from deepmd.pt_expt.long_range import ( + BaseLongRange, +) from deepmd.pt_expt.model.dp_zbl_model import ( DPZBLModel, ) @@ -75,6 +78,7 @@ pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, linear_atomic_model=LinearEnergyAtomicModel, + long_range_base=BaseLongRange, ) get_zbl_model = _model_factory.get_zbl_model @@ -92,8 +96,7 @@ def get_sezm_model(data: dict) -> BaseModel: ``bridging_method`` sugar expands to that form in :func:`get_model`); this builder rejects the flag. - Still unsupported here, each raising ``NotImplementedError``: the - virtual-atom (``deepspin``) spin scheme and ``lora``. + ``lora`` raises ``NotImplementedError``. Notes ----- @@ -107,11 +110,9 @@ def get_sezm_model(data: dict) -> BaseModel: """ data = copy.deepcopy(data) if "spin" in data: - if str(data["spin"].get("scheme", "deepspin")) != "native": + if str(data["spin"].get("scheme", "")) != "native": raise NotImplementedError( - "Spin DPA4/SeZM models with the virtual-atom (deepspin) " - "scheme are not supported in the pt_expt backend; use spin " - "scheme 'native' instead." + "DPA4/SeZM spin models require model.spin.scheme='native'." ) return get_native_spin_model(data) bridging_method = str(data.get("bridging_method", "none")) @@ -168,6 +169,7 @@ def get_sezm_model(data: dict) -> BaseModel: data["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types) descriptor, fitting, _ = _model_factory.get_model_components(data) + long_range = _model_factory.get_long_range(data, descriptor, fitting) return DPA4EnergyModel( descriptor=descriptor, fitting=fitting, @@ -175,6 +177,7 @@ def get_sezm_model(data: dict) -> BaseModel: atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=pair_exclude_types, preset_out_bias=data.get("preset_out_bias"), + **({} if long_range is None else {"long_range": long_range}), ) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 0dfd77ccc2..25e88487a7 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -23,6 +23,8 @@ from deepmd.dpmodel.utils.neighbor_graph import ( compact_nodes, expand_node_values, + flatten_topologies, + unflatten_topologies, ) from deepmd.pt_expt.common import ( auto_wrapped_class, @@ -43,6 +45,7 @@ from .edge_transform_output import ( fit_output_to_model_output_graph, + releaf_topologies, ) from .transform_output import ( fit_output_to_model_output, @@ -137,7 +140,8 @@ def _fused_energy_force_graph( """ am = model.atomic_model parts = am.fused_decomposition() - if parts is None: + if parts is None or parts[0].has_long_range(): + # The descriptor pipelines evaluate no long-range module. return None learned, pair_potential = parts desc, fit = learned.descriptor, learned.fitting_net @@ -191,6 +195,76 @@ def _fused_energy_force_graph( return ret +def check_owned_frames( + n_node: torch.Tensor, + n_local: torch.Tensor | None, + communicator: torch.Tensor | None = None, + periodic_marks: torch.Tensor | None = None, +) -> None: + """Reject halo nodes in the non-periodic frames of a single domain with a long-range module. + + The module sums over every pair of atoms of a non-periodic frame, so a + frame that one domain holds whole owns all its nodes; halo nodes would + carry no charge and remove their interactions from the sum. Across the + domains of a communicator, every domain owns its share of the frame + instead. The halo nodes of a periodic frame are images of its atoms, + which carry no charge: the reciprocal-space sum holds the images itself. + The counts are read eagerly only: a traced lower receives them from + callers that check them. + + Parameters + ---------- + n_node + Per-frame node counts with shape ``(nf,)``. + n_local + Per-frame owned node counts with shape ``(nf,)``, or ``None``. + communicator + The communicator input of the compact lower, a CPU tensor holding the + address of an MPI_Comm handle or zero, or ``None`` for a single domain. + periodic_marks + A tensor with a nonzero entry when the frames are periodic, such as + their cells or the frame mask of a reciprocal-space topology, or + ``None`` for non-periodic frames. It is read only after the fake + check, so that a traced lower computes nothing on it. + + Raises + ------ + ValueError + If a non-periodic frame of a single domain carries halo nodes. + """ + from torch._subclasses.fake_tensor import ( + is_fake, + ) + + if n_local is None or n_local is n_node or is_fake(n_local): + return + if communicator is not None and int(communicator.reshape(-1)[0]) != 0: + return + if periodic_marks is not None and bool(torch.any(periodic_marks != 0)): + return + if not torch.equal(n_local, n_node): + raise ValueError( + "a model with a long-range module evaluates complete non-periodic " + "frames on a single domain, whose graph owns every node (n_local " + f"== n_node), but got n_local {n_local.tolist()} for n_node " + f"{n_node.tolist()}" + ) + + +def _reciprocal_masks(topologies: dict[str, Any]) -> torch.Tensor | None: + """The mask of the reciprocal-space topology, true on periodic entries, or ``None``.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + ) + + masks = [ + topology.mask + for topology in topologies.values() + if not isinstance(topology, PairGraph) + ] + return masks[0] if masks else None + + def _pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor: """Append a single ``-1`` column to ``nlist`` for export-time tracing. @@ -356,7 +430,8 @@ class _WrapperForwardEnergyGraph: """Graph twin of :class:`_WrapperForwardEnergy` for the Hessian. Given flattened LOCAL coordinates for one frame, rebuilds the carry-all - ``NeighborGraph`` (so ``edge_vec`` tracks the coordinates through + ``NeighborGraph`` and the topologies the atomic model declares (so + ``edge_vec`` and the pair vectors track the coordinates through ``autograd.functional.hessian``'s double-backward) and returns the scalar reduced-output component. Unlike the dense wrapper this differentiates w.r.t. the ``(nloc, 3)`` LOCAL coordinates directly -- the carry-all graph @@ -399,6 +474,15 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: ng = build_neighbor_graph_for_method( self.method, cc, self.atype, self.box, self.rcut, self.pair_excl ) + # The topologies follow the same coordinates, so the Hessian holds + # every geometry path of the energy and the blocks coupling them. + topologies = self.model.atomic_model.build_topologies( + coord_flat.reshape(self.nloc, 3), + ng.n_node, + self.box, + atype=self.atype.reshape(-1), + pair_excl=self.pair_excl, + ) atomic_ret = self.model.atomic_model.forward_common_atomic_graph( ng, self.atype.reshape(-1), @@ -406,6 +490,7 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: aparam=self.aparam, spin=self.spin, charge_spin=self.charge_spin, + topologies=topologies, ) # atomic_ret[kk]: flat (N, *def), N == nloc for a single-frame carry-all # graph (all nodes owned); reduced output = sum over the node axis. @@ -608,6 +693,7 @@ def forward_common_lower_graph( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, torch.Tensor]: """Graph-native lower with autograd force/virial (dpa1/se_atten concat-tebd, attention included). @@ -696,6 +782,13 @@ def forward_common_lower_graph( DescrptBlockRepformers._exchange_ghosts_graph`). Only meaningful for descriptors whose ``has_message_passing_across_ranks()`` is ``True``. + topologies + The topologies the atomic model declares by + ``required_topologies()``, built on the node axis of the + graph. The geometry fields every topology names in its + ``LEAVES`` become autograd leaves next to ``edge_vec``, and + force and virial include their pullbacks. ``None`` (default) + for models declaring none. Returns ------- @@ -712,6 +805,17 @@ def forward_common_lower_graph( # make edge_vec the autograd leaf for the energy backward edge_vec = edge_vec.detach().requires_grad_(True) + if topologies: + if comm_dict is not None: + raise NotImplementedError( + "the long-range topologies describe complete systems; " + "a domain-decomposed system is not supported" + ) + check_owned_frames( + n_node, n_local, periodic_marks=_reciprocal_masks(topologies) + ) + # Every topology carries geometry leaves of its own. + topologies = releaf_topologies(topologies) if spin is not None: # second autograd leaf: force_mag = -dE/dspin (native spin). # Deliberately a SEPARATE leaf/backward from edge_vec rather @@ -751,6 +855,7 @@ def forward_common_lower_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, ) # ``forward_common_atomic_graph`` returns flat ``(N, *)`` output. # Pass directly to the flat-N transform; no rectangular reshape needed. @@ -772,6 +877,7 @@ def forward_common_lower_graph( # dynamic-edge torch.export. See fit_output_to_model_output_graph. node_capacity=atype.shape[0], n_local=n_local, + topologies=topologies, ) def _resolve_graph_method( @@ -892,6 +998,7 @@ def call_common_ragged( method = getattr(self, "neighbor_graph_method", None) if method is None: method = resolve_neighbor_graph_method("auto", coord.device) + pair_excl = getattr(self.atomic_model, "pair_excl", None) graph = build_ragged_neighbor_graph( method, coord, @@ -899,7 +1006,14 @@ def call_common_ragged( n_node, box, self.get_rcut(), - getattr(self.atomic_model, "pair_excl", None), + pair_excl, + ) + topologies = self.atomic_model.build_topologies( + coord.detach(), + graph.n_node, + box, + atype=atype, + pair_excl=pair_excl, ) predict = self.forward_common_lower_graph( atype, @@ -918,6 +1032,7 @@ def call_common_ragged( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) # The per-atom mask a rectangular batch carries is what tells the # loss each frame's real atom count. Nothing here is padded, so the @@ -1022,6 +1137,13 @@ def _call_common_graph( spin_flat = ( spin.reshape(n_padded, 3)[node_index] if spin is not None else None ) + topologies = self.atomic_model.build_topologies( + cc.detach().reshape(n_padded, 3)[node_index], + ng.n_node, + bb, + atype=atype_flat, + pair_excl=pair_excl, + ) model_predict = self.forward_common_lower_graph( atype_flat, ng.n_node, @@ -1039,6 +1161,7 @@ def _call_common_graph( aparam=ap_flat, spin=spin_flat, charge_spin=charge_spin, + topologies=topologies, ) # ``forward_common_lower_graph`` returns flat ``(N, *)`` per-atom # outputs over the real atoms. Scatter them back onto the padded @@ -1249,6 +1372,7 @@ def forward_common_lower_graph_exportable( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, destination_sorted: bool = False, + topologies: dict[str, Any] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """make_fx trace of ``forward_common_lower_graph`` with ``edge_vec`` @@ -1289,6 +1413,13 @@ def forward_common_lower_graph_exportable( output. ``None`` (default) is the existing, unconditioned trace with no mag output -- used by every non-spin caller of this generic (descriptor-agnostic) exportable. + topologies + Sample topologies the atomic model declares by + ``required_topologies()``, or ``None``. When given, they are + traced inputs laid out by + :func:`~deepmd.dpmodel.utils.neighbor_graph.flatten_topologies` + after ``charge_spin``; a model reading topologies carries no + native spin. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -1299,8 +1430,9 @@ def forward_common_lower_graph_exportable( A traced module whose ``forward`` accepts ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, fparam, aparam, charge_spin, spin)`` and - returns a dict with the same internal keys as + source_row_ptr, fparam, aparam, charge_spin, spin)`` -- or, + with topologies, ``spin`` replaced by the topology inputs -- + and returns a dict with the same internal keys as ``forward_common_lower_graph``. """ validate_graph_csr_for_export( @@ -1315,6 +1447,61 @@ def forward_common_lower_graph_exportable( ) model = self + if topologies: + topology_names = tuple(topologies) + + def fn_topologies( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + *topology_vals: torch.Tensor, + ) -> dict[str, torch.Tensor]: + return model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=destination_sorted, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + topologies=unflatten_topologies(topology_names, topology_vals), + ) + + return make_fx(fn_topologies, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + *flatten_topologies(topologies), + ) + # ``spin`` is a traced INPUT only when a real spin tensor is # given (native-spin models). For every non-spin caller # (``spin is None``) the traced ``fn`` must have the SAME arity as diff --git a/deepmd/pt_expt/model/native_spin_model.py b/deepmd/pt_expt/model/native_spin_model.py index 5103af1d03..1557f20fa4 100644 --- a/deepmd/pt_expt/model/native_spin_model.py +++ b/deepmd/pt_expt/model/native_spin_model.py @@ -329,10 +329,10 @@ def forward_lower_canonical_graph_exportable( has no conditional tail, so index 8 is the last slot. The magnetic force is a value here rather than an autograd result: the - fused operator emits the on-site half from the node kernel and the - neighbour half per edge, and the deployment path reduces the second - onto source nodes. ``mask_mag`` is derived from the types by the same - single owner the eager translation uses. + fused operator closes the on-site part in its node kernel and + accumulates the neighbour part of every edge onto the edge's source + node before it returns. ``mask_mag`` is derived from the types by the + same single owner the eager translation uses. Parameters ---------- diff --git a/deepmd/pt_expt/model/spin_model.py b/deepmd/pt_expt/model/spin_model.py index 83cc8ecba4..a35861b745 100644 --- a/deepmd/pt_expt/model/spin_model.py +++ b/deepmd/pt_expt/model/spin_model.py @@ -9,6 +9,9 @@ make_fx, ) +from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, +) from deepmd.dpmodel.model.spin_model import SpinModel as SpinModelDP from deepmd.pt_expt.common import ( torch_module, @@ -28,6 +31,12 @@ @torch_module class SpinModel(SpinModelDP): + def __init__(self, backbone_model: torch.nn.Module, spin: Spin) -> None: + """Initialize a spin wrapper around a compatible backbone.""" + if any(isinstance(module, DescrptDPA4) for module in backbone_model.modules()): + raise ValueError("DPA4/SeZM spin models require a native spin model.") + super().__init__(backbone_model=backbone_model, spin=spin) + def __getattr__(self, name: str) -> Any: """Get attribute from the wrapped model. diff --git a/deepmd/pt_expt/train/gradient.py b/deepmd/pt_expt/train/gradient.py index cc3b863015..7c3801ef46 100644 --- a/deepmd/pt_expt/train/gradient.py +++ b/deepmd/pt_expt/train/gradient.py @@ -24,6 +24,7 @@ ) import torch +import torch.distributed as dist if TYPE_CHECKING: from collections.abc import ( @@ -123,6 +124,8 @@ def update(self, total_norm: torch.Tensor) -> None: def raise_if_nonfinite( self, named_parameters: Callable[[], Iterable[tuple[str, torch.nn.Parameter]]], + *, + synchronize: bool = False, ) -> None: """ Raise if any norm accumulated since the previous call was non-finite. @@ -134,15 +137,22 @@ def raise_if_nonfinite( ---------- named_parameters : Callable[[], Iterable[tuple[str, torch.nn.Parameter]]] Accessor for the model's named parameters, consulted only when raising. + synchronize : bool, optional + Agree on divergence across all ranks before raising or entering + checkpoint collectives. Every rank must participate when enabled. Raises ------ RuntimeError If a non-finite gradient norm was recorded. """ - if self._nonfinite is None: - return - diverged = bool(self._nonfinite) + diverged = self._nonfinite is not None and bool(self._nonfinite) + if synchronize: + # Agree before diagnostics too: inspecting sharded gradients can + # itself require collectives, so all ranks must take that path. + gathered = [False] * dist.get_world_size() + dist.all_gather_object(gathered, diverged) + diverged = any(gathered) self._nonfinite = None if diverged: raise_nonfinite_gradient_norm(named_parameters()) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f1e7f27c58..5bcf798806 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -61,6 +61,15 @@ change_model_out_bias_by_task, resolve_step_schedule, ) +from deepmd.dpmodel.train.data import ( + add_training_data_requirements, +) +from deepmd.dpmodel.train.data_progress import ( + DATA_PROGRESS_CHECKPOINT_KEY, + collect_training_data_progress, + preserve_training_data_progress, + restore_training_data_progress, +) from deepmd.dpmodel.utils.batch import ( normalize_batch, split_batch, @@ -72,6 +81,10 @@ from deepmd.dpmodel.utils.learning_rate import ( make_learning_rate_schedule, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + flatten_topologies, + unflatten_topologies, +) from deepmd.dpmodel.utils.training_utils import ( compute_total_numb_batch, ) @@ -101,6 +114,14 @@ strip_saved_tensor_detach as _strip_saved_tensor_detach, ) from deepmd.pt.utils.compile_compat import trace_pad_dim as _trace_pad_dim +from deepmd.pt.utils.finetune import ( + _learned_part_path, + _move_learned_key, +) +from deepmd.pt.utils.spin_finetune import ( + compute_native_spin_input_stats, + get_native_spin_finetune_state, +) from deepmd.pt_expt.loss import ( DOSLoss, EnergyLoss, @@ -112,6 +133,9 @@ from deepmd.pt_expt.model import ( get_model, ) +from deepmd.pt_expt.model.edge_transform_output import ( + releaf_topologies, +) from deepmd.pt_expt.model.graph_lower import ( model_uses_graph_lower, ) @@ -146,9 +170,12 @@ from deepmd.pt_expt.utils.stat import ( make_stat_input, ) +from deepmd.utils import random as dp_random from deepmd.utils.data import ( DataRequirementItem, - min_pair_dist_requirement, +) +from deepmd.utils.data_group import ( + epoch_batch_counts, ) from deepmd.utils.data_system import ( DeepmdDataSystem, @@ -297,6 +324,10 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: atomic model and fitting. Shared weights do not make those branches interchangeable: tasks without a preset bias disable the reference, while a folded reference no longer appends isolated atoms to the graph. + + A long-range module is part of the traced computation and of its input + schema, so its identity (that of its first parameter) closes the key; + tasks share a graph only when they share the module as well. """ descriptor_id: int = 0 try: @@ -322,7 +353,17 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: break except AttributeError: pass - return (int(model.has_spin()), descriptor_id, fitting_id, *vacuum_state) + long_range_id = 0 + long_range = getattr(getattr(model, "atomic_model", None), "long_range", None) + if long_range is not None: + long_range_id = next((id(p) for p in long_range.parameters()), id(long_range)) + return ( + int(model.has_spin()), + descriptor_id, + fitting_id, + *vacuum_state, + long_range_id, + ) # --------------------------------------------------------------------------- @@ -765,6 +806,12 @@ def _trace_and_compile_graph( autograd leaf. Their compiled lower returns ``force_mag`` and ``mask_mag`` through the same model-owned translation used by eager execution. + A model declaring topologies (``atomic_model.required_topologies()``) + receives the fields of each, laid out by :func:`flatten_topologies`, as + further inputs ahead of the task buffers; their dynamic axes, the pairs, + the wave numbers or the axes of the mesh, are symbols distinct from the + other three. + Parameters ---------- model @@ -842,7 +889,9 @@ def _trace_and_compile_graph( # its graph builders omit the CSR metadata carried by deployment inputs. from deepmd.pt_expt.utils.serialization import ( build_synthetic_graph_inputs, + build_synthetic_topologies, check_graph_trace_torch_version, + check_mesh_compile_torch_version, count_synthetic_graph_edges, ) @@ -861,6 +910,7 @@ def _trace_and_compile_graph( # CUDA autograd-stream limitation) would otherwise mix a CPU model with # CUDA inputs and fail only on a GPU host. _trace_device = next(model.parameters()).device + check_mesh_compile_torch_version(model, _trace_device) e_real = count_synthetic_graph_edges( model, nframes=trace_nf, @@ -899,6 +949,22 @@ def _trace_and_compile_graph( # tensors, while forward_common_lower_graph takes it after them. s_spin = s_conditioning.pop(0) if spin is not None else None s_fparam, s_aparam, s_charge_spin = s_conditioning + # The dynamic axes of the topologies take sample lengths distinct from the + # other three and from the CSR offsets. + s_topologies = build_synthetic_topologies( + model, + trace_N, + _next_safe_prime( + e_max + 2, _forbidden | {trace_nf, trace_N, trace_N + 1, e_max} + ), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=_trace_device, + nframes=trace_nf, + forbidden=_forbidden | {e_max}, + ) + topology_names = tuple(s_topologies) + s_topology_vals = flatten_topologies(s_topologies) + n_topology_vals = len(s_topology_vals) def fn( atype: torch.Tensor, @@ -915,8 +981,10 @@ def fn( aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, spin: torch.Tensor | None, - *task_buf_vals: torch.Tensor, + *extra_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: + topologies = unflatten_topologies(topology_names, extra_vals[:n_topology_vals]) + task_buf_vals = extra_vals[n_topology_vals:] # Patch task-specific buffers with the proxy tensors so make_fx records # them as FX placeholders (mirrors the dense ``_trace_and_compile``). originals: dict[str, torch.Tensor | None] = {} @@ -950,6 +1018,7 @@ def fn( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) if spin is None: return _translate_energy_keys( @@ -996,6 +1065,7 @@ def fn( s_aparam, s_charge_spin, s_spin, + *s_topology_vals, *task_buf_vals_trace, ) @@ -1640,6 +1710,15 @@ def _forward_graph( aparam = aparam[node_index] if spin is not None: spin = spin[node_index] + # The topologies the atomic model declares are built here as well, + # outside the compiled region, on the node axis of the graph. + topologies = _model.atomic_model.build_topologies( + coord_3d if ragged else coord_3d.reshape(n_padded, 3)[node_index], + ng.n_node, + box_flat, + atype=atype_flat, + pair_excl=pair_excl, + ) # Lazy compile of the GRAPH lower (cached per structure key and mode). compiled_lower = self._compiled_lower_by_mode.get(self.training) @@ -1667,6 +1746,9 @@ def _forward_graph( edge_vec = ng.edge_vec.detach().requires_grad_(True) if spin is not None: spin = spin.detach().requires_grad_(True) + # The topologies enter as their fields, the geometry of every one a + # further leaf. + topology_vals = flatten_topologies(releaf_topologies(topologies)) if self._task_buf_order: try: @@ -1708,6 +1790,7 @@ def _forward_graph( aparam, charge_spin, spin, + *topology_vals, *task_buf_vals, ) self._report_pending_compile() @@ -1805,6 +1888,7 @@ def __init__( model_params = config["model"] training_params = config["training"] + self._configured_data_seed = training_params.get("seed") optimizer_params = config.get("optimizer", {}) validating_params = config.get("validating", {}) or {} @@ -1841,6 +1925,20 @@ def __init__( self.is_distributed = dist.is_available() and dist.is_initialized() self.rank = dist.get_rank() if self.is_distributed else 0 self.world_size = dist.get_world_size() if self.is_distributed else 1 + # Rank-local data draws must not change the shared task-selection sequence. + self._task_rng = np.random.default_rng( + None + if self._configured_data_seed is None + else self._configured_data_seed % (2**32) + ) + if ( + self.is_distributed + and self.multi_task + and self._configured_data_seed is None + ): + self._task_rng.bit_generator.state = self._broadcast_value_from_rank0( + self._task_rng.bit_generator.state if self.rank == 0 else None + ) self.sharding = ShardingPolicy.from_training_params( training_params, is_distributed=self.is_distributed ) @@ -1909,7 +2007,7 @@ def __init__( reset_epoch_streams() self.valid_numb_batch_by_task: dict[str, int] = {} - self.min_pair_dist_by_task: dict[str, float] = {} + self.pair_filter_enabled_by_task: dict[str, bool] = {} for model_key in self.model_keys: data_requirement = list(self.losses[model_key].label_requirement) data_requirement += get_additional_data_requirement(self.models[model_key]) @@ -1918,19 +2016,14 @@ def __init__( if self.multi_task else training_params ) - pair_filter = min_pair_dist_requirement( - self.model_params_by_task[model_key], - float( - (task_params.get("training_data") or {}).get("min_pair_dist", 0.0) - ), - ) - self.min_pair_dist_by_task[model_key] = ( - 0.0 if pair_filter is None else float(pair_filter.default) - ) - # Only training frames are filtered, so only the training data - # pays for the derived margin. - self.training_data_by_task[model_key].add_data_requirements( - data_requirement + ([] if pair_filter is None else [pair_filter]) + self.pair_filter_enabled_by_task[model_key] = ( + add_training_data_requirements( + self.training_data_by_task[model_key], + task_params["training_data"], + self.model_params_by_task[model_key], + data_requirement, + lambda data, requirements: data.add_data_requirements(requirements), + ) ) if self.validation_data_by_task[model_key] is not None: self.validation_data_by_task[model_key].add_data_requirements( @@ -2007,9 +2100,12 @@ def __init__( def _make_sample( _d: DeepmdDataSystem = _data, _n: int = _nbatch, - _min_pair_dist: float = self.min_pair_dist_by_task[model_key], + _filter_enabled: bool = self.pair_filter_enabled_by_task[model_key], ) -> list[dict[str, np.ndarray]]: - return make_stat_input(_d, _n, _min_pair_dist) + # Statistics run only on the chief. They must not advance its + # training cursors or data-source selection RNGs. + with preserve_training_data_progress({"statistics": _d}): + return make_stat_input(_d, _n, _filter_enabled) self._sample_funcs[model_key] = _make_sample @@ -2130,6 +2226,7 @@ def initialize_statistics( # cannot be copied into sharded parameters. ema_state_dict = None optimizer_state_dict = None + data_progress = None if resuming: log.info(f"Resuming from {resume_model}.") is_pte = resume_model.endswith((".pte", ".pt2")) @@ -2153,6 +2250,11 @@ def initialize_statistics( ema_state_dict = ( state_dict.get(EMA_CHECKPOINT_KEY) if continues_run else None ) + data_progress = ( + state_dict.get(DATA_PROGRESS_CHECKPOINT_KEY) + if continues_run + else None + ) state_dict = state_dict["model"] else: optimizer_state_dict = None @@ -2243,34 +2345,40 @@ def initialize_statistics( # Selective weight copy (per-branch key remapping) pretrained_state = pretrained_wrapper.state_dict() target_state = self._unwrapped.state_dict() - # A plain model and its zone-bridging composition hold the same - # learned descriptor and fitting net at different paths. - learned_paths = { - model_key: ( - _learned_part_path(self._unwrapped.model[model_key]), - _learned_part_path( - pretrained_wrapper.model[ - finetune_links[model_key].get_model_branch() - ] - ), - ) - for model_key in self.model_keys - } - new_state = {} - for key in target_state: - if key == "_extra_state": - new_state[key] = target_state[key] + new_state = target_state.copy() + # The atomic models, by module path, whose long-range module + # takes its initialization rather than pretrained weights. + fresh_long_range: dict[str, str] = {} + for model_key in self.model_keys: + finetune_rule = finetune_links[model_key] + source_key = finetune_rule.get_model_branch() + model = self._unwrapped.model[model_key] + source_model = pretrained_wrapper.model[source_key] + prefix = f"model.{model_key}." + if finetune_rule.convert_to_native_spin: + log.info( + "Initialize native-spin branch '%s' from spin-free branch '%s'.", + model_key, + source_key, + ) + converted = get_native_spin_finetune_state( + model, + source_model, + random_fitting=finetune_rule.get_random_fitting(), + ) + new_state.update( + {prefix + key: value for key, value in converted.items()} + ) continue - # Find which model_key this key belongs to - matched = False - for model_key in self.model_keys: - if f".{model_key}." not in key: + + # Plain and bridged models store their learned part at + # different paths; the analytical part keeps target state. + path = _learned_part_path(model) + source_path = _learned_part_path(source_model) + for key in target_state: + if not key.startswith(prefix): continue - matched = True - finetune_rule = finetune_links[model_key] - _key_from = finetune_rule.get_model_branch() - pretrained_key = key.replace(f".{model_key}.", f".{_key_from}.") - path, source_path = learned_paths[model_key] + pretrained_key = f"model.{source_key}." + key[len(prefix) :] if ( pretrained_key not in pretrained_state and path is not None @@ -2283,17 +2391,35 @@ def initialize_statistics( finetune_rule.get_random_fitting() and ".descriptor." not in key ) - if use_random: - new_state[key] = target_state[key] - elif pretrained_key in pretrained_state: - new_state[key] = pretrained_state[pretrained_key] + if use_random or pretrained_key not in pretrained_state: + if ".long_range." in key: + owner = key[: key.index(".long_range.")] + fresh_long_range[owner] = model_key else: - new_state[key] = target_state[key] - break - if not matched: - new_state[key] = target_state[key] + new_state[key] = pretrained_state[pretrained_key] self._unwrapped.load_state_dict(new_state) + # A long-range module added to the pretrained model takes the + # input statistics of the fine-tuning data, as in training from + # scratch, before the bias adjustment evaluates the model. + for owner, model_key in fresh_long_range.items(): + atomic_model = self._unwrapped.get_submodule(owner) + + def long_range_stat( + _atomic_model: torch.nn.Module = atomic_model, + _model_key: str = model_key, + ) -> None: + _atomic_model.compute_long_range_input_stat( + self._sample_funcs[_model_key] + ) + + self._run_stat_on_chief( + long_range_stat, + operation=f"long-range statistics for task {model_key!r}", + ) + if self.is_distributed: + self._broadcast_model_stat(atomic_model.long_range) + # Per-branch bias adjustment (rank 0 only, then broadcast) for model_key in self.model_keys: finetune_rule = finetune_links[model_key] @@ -2311,7 +2437,12 @@ def initialize_statistics( def update_finetune_bias( _model_key: str = model_key, _bias_mode: str = bias_mode, + _convert_to_native_spin: bool = finetune_rule.convert_to_native_spin, ) -> None: + if _convert_to_native_spin: + compute_native_spin_input_stats( + self.models[_model_key], self._sample_funcs[_model_key] + ) self.models[_model_key] = model_change_out_bias( self.models[_model_key], self._sample_funcs[_model_key], @@ -2464,6 +2595,12 @@ def update_finetune_bias( multi_task=self.multi_task, ) + # Data requirements, layout, statistics and compilation may all draw + # batches during setup. Restore only after those steps have finished. + self._restore_data_progress( + data_progress, reset=training_params.get("reset_data_progress", False) + ) + def _create_full_validators( self, *, @@ -2794,11 +2931,15 @@ def _next_training_batch(self, data_sys: Any, task_key: str) -> dict[str, Any]: RuntimeError If one local reader pass of consecutive batches holds no valid frame. """ - filter_enabled = self.min_pair_dist_by_task[task_key] > 0.0 + filter_enabled = self.pair_filter_enabled_by_task[task_key] discarded = 0 while True: - batch = normalize_batch(data_sys.get_batch()) - frame_mask = pair_margin_frame_mask(batch, filter_enabled) + raw_batch = data_sys.get_batch() + batch = normalize_batch(raw_batch) + batch_filter_enabled = getattr( + raw_batch, "pair_filter_enabled", filter_enabled + ) + frame_mask = pair_margin_frame_mask(batch, batch_filter_enabled) if frame_mask is None or frame_mask.all(): return batch if frame_mask.any(): @@ -2826,23 +2967,23 @@ def _epoch_length(self, model_key: str) -> int: Returns ------- int - Number of steps covering one pass over the task's training data. + Nominal number of optimizer steps in one training epoch. Notes ----- - A data system reports ``nbatches[i]``, the batch count of system ``i``, - and ``sys_probs[i]``, the probability of drawing from that system, from - which ``compute_total_numb_batch`` derives the dataset-wide epoch - length ``ceil(max_i(nbatches[i] / sys_probs[i]))``. LMDB data reports - that global count while its sampler shards batches evenly across ranks; - legacy data systems remain replicated. In both cases one rank takes - ``ceil(total / world_size)`` optimizer steps per scheduled epoch. This - differs from the local reader pass length for replicated data, which - the data system exposes through ``get_batch_pass_length``. + The epoch length is ``ceil(max_i(local_batches[i] / sys_probs[i]))``. + Each LMDB contributes its sampler's rank-local workload, including + its own distributed padding. Replicated NumPy readers contribute + their global batch counts divided by world size. Size-based sampling + uses global source sizes without per-rank padding; uniform sampling + corrects for each native stream's rank-averaged batch frame count. + This scheduled workload differs from a replicated reader's pass length, + which is exposed separately through ``get_batch_pass_length``. """ data = self.training_data_by_task[model_key] - total = compute_total_numb_batch(data.nbatches, data.sys_probs) - return int(np.ceil(total / self.world_size)) + return compute_total_numb_batch( + epoch_batch_counts(data, self.world_size), data.sys_probs + ) # ------------------------------------------------------------------ # Distribution helpers @@ -3035,6 +3176,7 @@ def _broadcast_model_stat(model: torch.nn.Module) -> None: def _broadcast_value_from_rank0(self, value: Any) -> Any: """Return rank 0's copy of ``value`` on every rank. + Unseeded task-selection RNG state must agree across ranks at startup. Epoch lengths round a quotient of sampling probabilities that is often an exact integer in real arithmetic, so a last-bit difference between ranks -- as reduction kernels dispatched for different CPU features @@ -3055,15 +3197,100 @@ def _broadcast_value_from_rank0(self, value: Any) -> Any: @property def checkpoint_is_collective(self) -> bool: - """Whether assembling a checkpoint needs every rank.""" - return self.sharding.enabled + """Whether checkpoint preflight and data collection need every rank.""" + # Frame filtering can consume different numbers of local batches, so + # even plain DDP must contribute every rank's cursor at save time. + return self.is_distributed or self.sharding.enabled + + def _collect_data_progress(self) -> dict[str, Any]: + """Gather compact rank-local progress only at a checkpoint boundary.""" + local_progress = collect_training_data_progress(self.training_data_by_task) + ranks: list[Any] = [local_progress] + if self.is_distributed: + ranks = [None] * self.world_size + dist.all_gather_object(ranks, local_progress) + return { + "version": 1, + "world_size": self.world_size, + "configured_seed": self._configured_data_seed, + "task_rng": self._task_rng.bit_generator.state, + "ranks": ranks, + } + + def _restore_data_progress( + self, progress: Mapping[str, Any] | None, *, reset: bool = False + ) -> None: + """Restore progress unless an explicit reset or topology change prevents it. + + Validate the rank envelope before either fallback. An explicit reset + retains all newly initialized sources and the selection RNG together, + without partially applying saved task state or bypassing model restart. + """ + if progress is None: + return + if not isinstance(progress, Mapping): + raise ValueError("Data-progress checkpoint must be a mapping.") + if type(progress.get("version")) is not int or progress["version"] != 1: + raise ValueError("Unsupported distributed data-progress version.") + saved_world_size = progress.get("world_size") + if type(saved_world_size) is not int or saved_world_size < 1: + raise ValueError("Data-progress checkpoint world size must be positive.") + ranks = progress.get("ranks") + if not isinstance(ranks, list) or len(ranks) != saved_world_size: + raise ValueError("Data-progress checkpoint must contain every rank.") + if not all(isinstance(state, Mapping) for state in ranks): + raise ValueError("Data-progress checkpoint rank states must be mappings.") + if reset: + if self.rank == 0: + log.warning( + "training.reset_data_progress is true. Model, optimizer and " + "step are restored, but training data order and selection " + "RNG restart from their newly initialized state. Batches " + "may repeat; exact data-order continuation is unavailable." + ) + return + if saved_world_size != self.world_size: + # Model/optimizer/step restart has always allowed a new topology. + # Rank-local cursors cannot be remapped safely, particularly after + # filtering, so retain the freshly initialized data and RNG instead. + if self.rank == 0: + log.warning( + "Checkpoint world size changed from %d to %d. Model, optimizer " + "and step are restored, but training data order and selection " + "RNG restart from their newly initialized state; exact " + "data-order continuation is unavailable.", + saved_world_size, + self.world_size, + ) + return + if "configured_seed" not in progress or "task_rng" not in progress: + raise ValueError( + "Data-progress checkpoint lacks its seed or task RNG identity; " + "set training.reset_data_progress=true to restart data sampling." + ) + if progress["configured_seed"] != self._configured_data_seed: + raise ValueError( + "Training data progress is incompatible: training.seed changed; " + "set training.reset_data_progress=true to restart data sampling." + ) + task_rng = np.random.default_rng(0) + task_rng.bit_generator.state = progress["task_rng"] + restore_training_data_progress(self.training_data_by_task, ranks[self.rank]) + self._task_rng = task_rng + + def ensure_finite_gradients_for_checkpoint(self) -> None: + """Check the accumulated gradient norm on all ranks before publication.""" + self.nonfinite_grad_guard.raise_if_nonfinite( + self.wrapper.named_parameters, synchronize=self.is_distributed + ) def save_checkpoint(self, step: int) -> None: - # Abort before writing if any gradient norm since the previous - # checkpoint was non-finite, so a diverged interval is not persisted. - self.nonfinite_grad_guard.raise_if_nonfinite(self.wrapper.named_parameters) + self.ensure_finite_gradients_for_checkpoint() + data_progress = self._collect_data_progress() + if self.rank != 0 and not self.sharding.enabled: + return ckpt_path = self.ckpt_store.path_for(step) - self._save_checkpoint_to_path(ckpt_path, step=step) + self._save_checkpoint_to_path(ckpt_path, step=step, data_progress=data_progress) if self.rank == 0: self.ckpt_store.publish(ckpt_path) self.ckpt_store.prune(ckpt_path) @@ -3080,10 +3307,12 @@ def _save_full_validation_checkpoint( save_path: Path, lr: float = 0.0, step: int = 0, + *, + data_progress: dict[str, Any] | None = None, ) -> None: """Save a checkpoint selected by full validation.""" del lr - self._save_checkpoint_to_path(save_path, step=step) + self._save_checkpoint_to_path(save_path, step=step, data_progress=data_progress) def _save_full_validation_ema_checkpoint( self, @@ -3105,6 +3334,7 @@ def _save_checkpoint_to_path( *, step: int, use_ema_weights: bool = False, + data_progress: dict[str, Any] | None = None, ) -> None: """Serialize the current trainer state to an explicit checkpoint path. @@ -3119,6 +3349,9 @@ def _save_checkpoint_to_path( Such a checkpoint is a deployment snapshot: it carries neither the optimizer state nor the EMA state, both of which describe the live weights it does not contain. + data_progress : dict or None, optional + Rank-local progress already gathered before chief-only validation. + Ordinary checkpoints collect it on every rank while being written. """ if use_ema_weights: with self.model_ema.apply_shadow(self.model): @@ -3129,7 +3362,7 @@ def _save_checkpoint_to_path( include_ema_state=False, ) return - self._write_checkpoint(ckpt_path, step=step) + self._write_checkpoint(ckpt_path, step=step, data_progress=data_progress) def _write_checkpoint( self, @@ -3138,6 +3371,7 @@ def _write_checkpoint( step: int, include_optimizer: bool = True, include_ema_state: bool = True, + data_progress: dict[str, Any] | None = None, ) -> None: """Serialize the wrapper, and optionally optimizer and EMA state.""" self._unwrapped.train_infos["step"] = step @@ -3160,6 +3394,8 @@ def _write_checkpoint( finally: for task_key, compiled in compiled_backup.items(): wrapper.model[task_key] = compiled + if include_optimizer and data_progress is None: + data_progress = self._collect_data_progress() # Sharded state is assembled on the chief; the other ranks have played # their part in the collectives above and hold nothing to write. if self.rank != 0: @@ -3167,6 +3403,7 @@ def _write_checkpoint( state: dict[str, Any] = {"model": model_state} if optim_state is not None: state["optimizer"] = optim_state + state[DATA_PROGRESS_CHECKPOINT_KEY] = data_progress if include_ema_state and self.model_ema is not None: state[EMA_CHECKPOINT_KEY] = self.model_ema.state_dict() ckpt_path.parent.mkdir(parents=True, exist_ok=True) @@ -3276,15 +3513,16 @@ def _precompile_outside_collectives(self) -> None: trainable_parameters = tuple( parameter for parameter in inner.parameters() if parameter.requires_grad ) - for task in self.training_tasks: - input_dict, label_dict = self.get_data(is_train=True, task_key=task.key) - _, loss, _ = inner( - **input_dict, - cur_lr=self.scheduler.get_last_lr()[0], - label=label_dict, - task_key=task.key, - ) - torch.autograd.grad(loss, trainable_parameters, allow_unused=True) + with preserve_training_data_progress(self.training_data_by_task): + for task in self.training_tasks: + input_dict, label_dict = self.get_data(is_train=True, task_key=task.key) + _, loss, _ = inner( + **input_dict, + cur_lr=self.scheduler.get_last_lr()[0], + label=label_dict, + task_key=task.key, + ) + torch.autograd.grad(loss, trainable_parameters, allow_unused=True) if torch.cuda.is_available(): torch.cuda.synchronize() log.info( @@ -3309,8 +3547,7 @@ def run(self) -> None: super().run(self.training_tasks) if self.change_bias_after_training and self.num_steps > self.start_step: self._change_bias_after_training() - if self.rank_context.is_chief: - self.save_checkpoint(self.num_steps) + self.save_checkpoint(self.num_steps) finally: self._close_data_systems() if self.rank_context.is_chief: @@ -3352,26 +3589,45 @@ def run_full_validation( learning_rate: float, ) -> None: """Run the active full validation flows for one step.""" - for validator, save_checkpoint in ( - (self.full_validator, self._save_full_validation_checkpoint), - (self.ema_full_validator, self._save_full_validation_ema_checkpoint), + for validator, save_checkpoint, include_progress in ( + (self.full_validator, self._save_full_validation_checkpoint, True), + (self.ema_full_validator, self._save_full_validation_ema_checkpoint, False), ): - if validator is None: + if validator is None or not validator.should_run(display_step): continue - validator.run( - step_id=display_step, - display_step=display_step, - lr=learning_rate, - save_checkpoint=save_checkpoint, - ) + # Plain DDP's validator calls its save hook only on the chief. + # Gather before entering it, while every rank is here, and pass + # the result through so that hook never starts another collective. + if include_progress and validator.save_best: + save_checkpoint = functools.partial( + save_checkpoint, data_progress=self._collect_data_progress() + ) + random_state = dp_random.get_state() + + def save_with_training_rng( + *args: Any, _save: Callable = save_checkpoint, **kwargs: Any + ) -> None: + # Best checkpoints are written inside validation. Record the + # training RNG, not the draws made by validation-data reads. + dp_random.set_state(random_state) + _save(*args, **kwargs) + + try: + validator.run( + step_id=display_step, + display_step=display_step, + lr=learning_rate, + save_checkpoint=save_with_training_rng, + prepare_checkpoint=self.ensure_finite_gradients_for_checkpoint, + ) + finally: + dp_random.set_state(random_state) def select_task(self, tasks: TrainingTaskCollection) -> TrainingTask: - """Select a task using DeePMD's seeded random helper.""" + """Select a task independently of rank-local data draws.""" if not tasks.is_multitask: return tasks[tasks.keys[0]] - from deepmd.utils import random as dp_random - - model_index = dp_random.choice( + model_index = self._task_rng.choice( np.arange(len(tasks), dtype=np.int_), p=tasks.probabilities, ) @@ -3393,12 +3649,22 @@ def collect_display_results( self.wrapper.eval() self._display_cur_lr_sched = step_result.payload["cur_lr_sched"] try: - return super().collect_display_results( - tasks, - active_task=active_task, - step=step, - step_result=step_result, - ) + # Only inactive tasks need extra training batches for display. + # The chief alone evaluates them, so those reads must be peeks. + # Preserving the RNG also isolates validation, whose cursor is not + # part of a training checkpoint, from future task/system choices. + display_data = { + task.key: task.training_data + for task in tasks + if task.key != active_task.key and self.metric_accumulator is None + } + with preserve_training_data_progress(display_data): + return super().collect_display_results( + tasks, + active_task=active_task, + step=step, + step_result=step_result, + ) finally: self._display_cur_lr_sched = None self.wrapper.train() @@ -3540,37 +3806,11 @@ def _more_loss_to_float(cls, more_loss: dict[str, Any]) -> dict[str, float]: } -def _learned_part_path(model: torch.nn.Module) -> str | None: - """Return the state-dict path of the learned descriptor-fitting part of a model. - - A plain model holds that part as its atomic model and a composition with - one learned sub-model, such as a zone-bridging one, as that sub-model, so - fine-tuning between the two transfers the part across the paths. A model - without a single learned part returns ``None``. - """ - atomic_model = model.atomic_model - if getattr(atomic_model, "descriptor", None) is not None: - return "atomic_model." - learned = [ - index - for index, child in enumerate(getattr(atomic_model, "models", ())) - if getattr(child, "descriptor", None) is not None - ] - return f"atomic_model.models.{learned[0]}." if len(learned) == 1 else None - - -def _move_learned_key(key: str, path: str, source_path: str) -> str: - """Rewrite a descriptor or fitting-net key from one learned-part path to another.""" - for part in ("descriptor.", "fitting_net."): - if f".{path}{part}" in key: - return key.replace(f".{path}{part}", f".{source_path}{part}", 1) - return key - - def model_change_out_bias( _model: Any, _sample_func: Any, _bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, ) -> Any: """Change the output bias of a model based on sampled data. @@ -3582,6 +3822,9 @@ def model_change_out_bias( Callable that returns sampled data for bias computation. _bias_adjust_mode ``"change-by-statistic"`` or ``"set-by-statistic"``. + preset_out_bias + Optional preset override. An empty mapping ignores checkpoint presets; + omitting it preserves training and fine-tuning behavior. Returns ------- @@ -3596,6 +3839,7 @@ def model_change_out_bias( _sample_func, bias_adjust_mode=_bias_adjust_mode, recompute_input_stats=isinstance(_model, DPModelCommon), + preset_out_bias=preset_out_bias, ) diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py index 746b6df91a..b63f6d1fbf 100644 --- a/deepmd/pt_expt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -14,6 +14,7 @@ ) from dataclasses import ( dataclass, + replace, ) from pathlib import ( Path, @@ -56,6 +57,9 @@ normalize_full_validation_metric, resolve_full_validation_start_step, ) +from deepmd.utils.data_group import ( + iter_data_sources, +) from deepmd.utils.eval_metrics import ( ENERGY_FULL_VALIDATION_PROFILE, SPIN_FULL_VALIDATION_PROFILE, @@ -294,11 +298,11 @@ def __init__( self.checkpoint_dir.mkdir(parents=True, exist_ok=True) self._initialize_best_checkpoints(restart_training=restart_training) - # Lazily-populated full test snapshot for LMDB validation. Frames are + # Lazily-populated full test snapshots for LMDB validation. Frames are # evaluated in atom-count and label-availability groups so each stack # has consistent shapes and scalar find_* flags. Reuse the decoded - # snapshot on subsequent validation calls. - self._lmdb_test_data: LmdbTestData | None = None + # snapshots on subsequent validation calls, independently for each source. + self._lmdb_snapshots: dict[int, LmdbTestData] = {} def should_run(self, display_step: int) -> bool: """Check whether the current step should trigger full validation.""" @@ -315,8 +319,14 @@ def run( display_step: int, lr: float, save_checkpoint: Any, + prepare_checkpoint: Callable[[], None] | None = None, ) -> FullValidationResult | None: - """Run full validation if the current step is due.""" + """Run full validation if the current step is due. + + ``prepare_checkpoint`` runs on every rank only for a selected checkpoint. + It must agree on distributed failures before returning or raising, so + no peer enters checkpoint collectives after another rejects the save. + """ if not self.should_run(display_step): return None @@ -327,10 +337,15 @@ def run( caught_exception: Exception | None = None error_message = None save_path = [None] + proposed_records = self.topk_records if self.rank == 0: try: result = self._evaluate(display_step) - save_path[0] = result.saved_best_path + proposed_records, save_path[0] = self._propose_best_state( + display_step=display_step, + selected_metric_value=result.selected_metric_value, + ) + result = replace(result, saved_best_path=save_path[0]) except Exception as exc: caught_exception = exc error_message = ( @@ -343,6 +358,15 @@ def run( if self.is_distributed: dist.broadcast_object_list(save_path, src=0) + if save_path[0] is not None and prepare_checkpoint is not None: + # All ranks agree before mutating top-K metadata or entering a + # serializer that may gather sharded model/optimizer state. + prepare_checkpoint() + + if self.rank == 0: + self.topk_records = proposed_records + self._sync_state_store() + if save_path[0] is not None: try: # Assembling a checkpoint from shards is collective, so every @@ -384,7 +408,7 @@ def run( return result if self.rank == 0 else None def _evaluate(self, display_step: int) -> FullValidationResult: - """Evaluate all validation systems and update best state.""" + """Evaluate all validation systems without publishing checkpoint state.""" # === Step 1. Switch to Evaluation Mode === was_training = bool(getattr(self.model, "training", True)) self.model.eval() @@ -408,18 +432,13 @@ def _evaluate(self, display_step: int) -> FullValidationResult: f"validation dataset: {self.metric_name.upper()}." ) - # === Step 3. Update Best Tracking === selected_metric_value = float(metrics[self.metric_key]) - saved_best_path = self._update_best_state( - display_step=display_step, - selected_metric_value=selected_metric_value, - ) return FullValidationResult( display_step=display_step, metrics=metrics, selected_metric_key=self.metric_key, selected_metric_value=selected_metric_value, - saved_best_path=saved_best_path, + saved_best_path=None, ) def evaluate_all_systems(self) -> dict[str, float]: @@ -457,28 +476,27 @@ def _iter_validation_data_systems(self) -> Iterator[Any]: - A loader set owns ``systems``, each wrapping a ``DeepmdData`` in ``data_system``. """ - validation_data = self.validation_data - if hasattr(validation_data, "_reader"): - lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data) - for nloc in sorted(lmdb_test_data.nloc_groups): - for indices in lmdb_test_data.availability_groups( - lmdb_test_data.nloc_groups[nloc] - ): - yield LmdbTestDataNlocView(lmdb_test_data, nloc, indices) - return - - if hasattr(validation_data, "data_systems"): - yield from validation_data.data_systems - return - - for dataset in validation_data.systems: - data_system = getattr(dataset, "data_system", None) - if data_system is None: - raise TypeError( - "Full validation expects each dataset in validation_data.systems " - f"to expose a `data_system`, got {type(dataset)!r}." + for source in iter_data_sources(self.validation_data): + if hasattr(source, "_reader"): + snapshot = self._get_lmdb_test_data_snapshot(source) + yield from ( + LmdbTestDataNlocView(snapshot, nloc, indices) + for nloc in sorted(snapshot.nloc_groups) + for indices in snapshot.availability_groups( + snapshot.nloc_groups[nloc] + ) ) - yield data_system + elif hasattr(source, "data_systems"): + yield from source.data_systems + else: + for dataset in source.systems: + data_system = getattr(dataset, "data_system", None) + if data_system is None: + raise TypeError( + "Full validation expects each dataset in validation_data.systems " + f"to expose a `data_system`, got {type(dataset)!r}." + ) + yield data_system def _get_lmdb_test_data_snapshot(self, lmdb_dataset: Any) -> LmdbTestData: """Build (once) and return the cached LMDB test snapshot. @@ -488,8 +506,9 @@ def _get_lmdb_test_data_snapshot(self, lmdb_dataset: Any) -> LmdbTestData: the full-validation snapshot sees exactly the same fields and dtypes as training batches. """ - if self._lmdb_test_data is not None: - return self._lmdb_test_data + key = id(lmdb_dataset) + if key in self._lmdb_snapshots: + return self._lmdb_snapshots[key] reader = getattr(lmdb_dataset, "_reader", None) lmdb_path = getattr(lmdb_dataset, "lmdb_path", None) @@ -512,14 +531,15 @@ def _get_lmdb_test_data_snapshot(self, lmdb_dataset: Any) -> LmdbTestData: f"{type(lmdb_dataset)!r}." ) - self._lmdb_test_data = LmdbTestData( + snapshot = LmdbTestData( lmdb_path, type_map=list(type_map), shuffle_test=False, ) if data_requirements: - self._lmdb_test_data.add_data_requirement(data_requirements) - return self._lmdb_test_data + snapshot.add_data_requirement(data_requirements) + self._lmdb_snapshots[key] = snapshot + return snapshot def _evaluate_system( self, data_system: DeepmdData @@ -696,13 +716,13 @@ def predict_batch( prediction["virial"] = batch_prediction["virial"] return prediction - def _update_best_state( + def _propose_best_state( self, *, display_step: int, selected_metric_value: float, - ) -> str | None: - """Update the top-K records and return the checkpoint path to save.""" + ) -> tuple[list[BestCheckpointRecord], str | None]: + """Propose top-K records and a save path without changing published state.""" candidate = BestCheckpointRecord( metric=selected_metric_value, step=display_step, @@ -714,14 +734,17 @@ def _update_best_state( updated_records.sort() updated_records = updated_records[: self.max_best_ckpt] if candidate not in updated_records: - return None - - self.topk_records = updated_records - self._sync_state_store() - if not self.save_best: - return None - candidate_rank = self.topk_records.index(candidate) + 1 - return str(self._best_checkpoint_path(display_step, candidate_rank)) + return self.topk_records, None + save_path = ( + str( + self._best_checkpoint_path( + display_step, updated_records.index(candidate) + 1 + ) + ) + if self.save_best + else None + ) + return updated_records, save_path def _sync_state_store(self) -> None: """Synchronize top-K validation state into the configured state store.""" diff --git a/deepmd/pt_expt/utils/auto_neighbor_list.py b/deepmd/pt_expt/utils/auto_neighbor_list.py new file mode 100644 index 0000000000..c0ed26344d --- /dev/null +++ b/deepmd/pt_expt/utils/auto_neighbor_list.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Batch-aware selection of persistent torch neighbor-list strategies.""" + +import warnings +from typing import ( + Any, + Literal, +) + +import torch + +from deepmd.dpmodel.utils.default_neighbor_list import ( + DefaultNeighborList, +) +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, +) +from deepmd.dpmodel.utils.neighbor_list import ( + EdgeNeighborList, + NeighborList, +) +from deepmd.pt.utils.nv_nlist import ( + NvNeighborList, + is_nv_available, +) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + VesinNeighborList, + is_vesin_torch_available, +) + + +class _NativeNeighborList(DefaultNeighborList): + """Build native extended inputs or uncapped edge-vector inputs.""" + + def build( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + sel: list[int], + return_mode: str = "extended", + pair_excl: PairExcludeMask | None = None, + ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: + if return_mode != "edges": + return super().build(coord, atype, box, rcut, sel, return_mode, pair_excl) + warnings.warn( + "Neither vesin[torch] nor nvalchemi-toolkit-ops is available; " + "using dense O(N^2) neighbor search (including any periodic images), " + "which can exhaust memory for large systems. Install vesin[torch]>=0.6 " + "or nvalchemi-toolkit-ops to enable cell-list neighbor search.", + RuntimeWarning, + stacklevel=2, + ) + graph = build_neighbor_graph(coord, atype, box, rcut, pair_excl=pair_excl) + return EdgeNeighborList( + coord=coord.reshape(*atype.shape, 3).contiguous(), + atype=atype.contiguous(), + edge_index=graph.edge_index, + edge_vec=graph.edge_vec, + edge_scatter_index=graph.edge_index, + edge_mask=graph.edge_mask & (graph.edge_vec.square().sum(dim=-1) > 1e-10), + ) + + +class AdaptiveNeighborList(NeighborList): + """Select a cell-list backend from the device and batch shape. + + ``auto`` favors persistent Vesin searches for single frames and Toolkit-Ops + for multi-frame batches on CPU and CUDA. ``native`` retains the + unbuffered batched policy: Vesin for one CPU frame, Toolkit-Ops otherwise, + with either backend serving as the fallback when the other is unavailable. + Without either optional backend, native search preserves the requested + representation: edge inputs use an uncapped dense graph, not a capped nlist. + The selection is independent of the requested neighbor representation. + """ + + def __init__( + self, skin: float = 0.0, *, method: Literal["auto", "native"] = "auto" + ) -> None: + self._method = method + self._vesin = VesinNeighborList(skin=0.0 if method == "native" else skin) + self._nv = NvNeighborList() + self._native = _NativeNeighborList() + + def build( + self, + coord: Any, + atype: Any, + box: Any, + rcut: float, + sel: list[int], + return_mode: str = "extended", + pair_excl: PairExcludeMask | None = None, + ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: + """Build the requested schema with the available inference strategy.""" + vesin_available = is_vesin_torch_available() + single_vesin = atype.shape[0] == 1 and ( + self._method == "auto" or coord.device.type == "cpu" + ) + if single_vesin and vesin_available: + builder = self._vesin + elif is_nv_available(): + builder = self._nv + elif vesin_available: + builder = self._vesin + else: + builder = self._native + return builder.build( + coord, atype, box, rcut, sel, return_mode=return_mode, pair_excl=pair_excl + ) + + +def make_neighbor_list( + method: str, + device: torch.device, + skin: float = 0.0, + *, + return_mode: Literal["extended", "edges"] = "extended", +) -> NeighborList | None: + """Select a search strategy compatible with the lower-input contract. + + ``None`` delegates extended-coordinate construction to the caller's native + builder. Edge inputs require an explicit carry-all builder: falling back to + a fixed-capacity neighbor list would discard neighbors before conversion. + """ + if method == "native": + if return_mode == "edges": + return AdaptiveNeighborList(method="native") + return None + if method == "auto": + if ( + return_mode == "edges" + or is_vesin_torch_available() + or (device.type in {"cpu", "cuda"} and is_nv_available()) + ): + return AdaptiveNeighborList(skin=skin) + return None + if method == "vesin": + if not is_vesin_torch_available(): + raise ImportError( + "Vesin neighbor search requires vesin[torch]>=0.6; " + "install it or select the native neighbor builder." + ) + return VesinNeighborList(skin=skin) + if method == "nv": + if device.type not in {"cpu", "cuda"}: + raise ValueError("NV neighbor search requires CPU or CUDA tensors") + if not is_nv_available(): + raise ImportError("NV neighbor search requires nvalchemi-toolkit-ops") + return NvNeighborList() + raise ValueError( + f"Unknown neighbor backend {method!r}; expected 'auto', 'native', 'vesin', or 'nv'." + ) diff --git a/deepmd/pt_expt/utils/canonical_graph.py b/deepmd/pt_expt/utils/canonical_graph.py index e01163a073..c8afd7d63a 100644 --- a/deepmd/pt_expt/utils/canonical_graph.py +++ b/deepmd/pt_expt/utils/canonical_graph.py @@ -19,16 +19,51 @@ NeighborGraph, ) -UINT32_MAX = (1 << 32) - 1 +INT32_MAX = (1 << 31) - 1 + + +def canonical_index_dtype(node_count: int) -> torch.dtype: + """Choose node-index width; edge counts and CSR positions always use int64. + + Four-byte words retain a signed representation at native host boundaries. + """ + return torch.uint32 if node_count <= INT32_MAX + 1 else torch.int64 + + +def canonical_index_bytes(index: torch.Tensor) -> torch.Tensor: + """Expose contiguous integer words as bytes with a final word-size axis.""" + return index.view(torch.uint8).reshape(*index.shape, index.element_size()) + + +def canonical_index_words(storage: torch.Tensor) -> torch.Tensor: + """Restore a canonical byte carrier to integer words without copying.""" + if storage.dtype != torch.uint8: + return storage + if storage.ndim < 2 or not storage.is_contiguous(): + raise ValueError( + "canonical index bytes require contiguous storage and a word axis" + ) + width = storage.shape[-1] + if width not in (4, 8): + raise ValueError( + f"canonical index word width must be 4 or 8 bytes, got {width}" + ) + dtype = torch.int64 if width == 8 else torch.uint32 + return storage.view(dtype).squeeze(-1) @dataclass(frozen=True) class CanonicalGraph: """Store a cutoff-compacted destination-major graph without redundant fields. - The physical edge count is the final value of either CSR row-pointer tensor. - The storage edge axis has length ``max(E, 2)``; suffix entries are guards - excluded from every CSR row. + The physical edge count is the final value of the destination row-pointer + tensor. The storage edge axis has length at least ``max(E, 2)``; entries + past the physical edges are guards excluded from every CSR row. + + The source CSR serves the consumers that group edges by source: the DPA1 + force assembly and the source-attributed per-atom virial of DPA4C. A DPA4C + graph whose caller declines the per-atom virial may carry two-entry + placeholders in its place. Parameters ---------- @@ -38,16 +73,18 @@ class CanonicalGraph: Per-frame owned node counts with shape ``(nf,)``, int64. source Source-node index for each edge storage slot with shape ``(S,)``, - uint32. + uint32 or int64, selected from the node extent. + Exported programs carry these words as uint8 with shape ``(S, W)``, + where the dynamic word width ``W`` is four or eight bytes. edge_vec Neighbor-minus-center vectors with shape ``(S, 3)``, float32. destination_row_ptr Destination CSR offsets with shape ``(N + 1,)``, int64. source_row_ptr - Source CSR offsets with shape ``(N + 1,)``, int64. + Source CSR offsets with shape ``(N + 1,)``, int64, or a placeholder. source_order - Edge storage positions grouped by source with shape ``(S,)`` and the - same dtype as ``source``. + Edge storage positions grouped by source with shape ``(S,)``, int64, + or a placeholder. This width is independent of the node indices. """ n_node: torch.Tensor @@ -63,12 +100,20 @@ def validate_canonical_graph_shapes( graph: CanonicalGraph, node_count: int, ) -> None: - """Validate shape, dtype, and device invariants without reading tensor data.""" + """Validate shape, dtype, and device invariants without reading tensor data. + + The extents of the source CSR are left to the operators that read it, + because a placeholder is a valid source CSR for a DPA4C graph and the + deployment trace keeps its extents independent of the node and edge axes. + """ index_dtype = graph.source.dtype - if index_dtype != torch.uint32: - raise ValueError("canonical graph source must be uint32") - if graph.source_order.dtype != index_dtype: - raise ValueError("canonical graph source and source_order dtypes must match") + byte_storage = index_dtype == torch.uint8 + if index_dtype not in (torch.uint8, torch.uint32, torch.int64): + raise ValueError( + "canonical graph source must be uint32, int64 or exported bytes" + ) + if graph.source_order.dtype != torch.int64: + raise ValueError("canonical graph source_order must be int64") if graph.edge_vec.dtype != torch.float32: raise ValueError("canonical graph edge_vec must be float32") if graph.n_node.dtype != torch.int64 or graph.n_local.dtype != torch.int64: @@ -78,31 +123,24 @@ def validate_canonical_graph_shapes( or graph.source_row_ptr.dtype != torch.int64 ): raise ValueError("canonical graph row pointers must be int64") - if graph.source.ndim != 1 or graph.source_order.shape != graph.source.shape: - raise ValueError("canonical graph source arrays must share shape (S,)") + if graph.source.ndim != (2 if byte_storage else 1) or graph.source_order.ndim != 1: + raise ValueError("canonical graph source arrays have invalid rank") + # A traced word width is a symbol of range [4, 8]: a bound adds no guard, + # and the operators that read the words accept only four or eight bytes. + if byte_storage and not 4 <= graph.source.shape[-1] <= 8: + raise ValueError("canonical index word width must be four or eight bytes") + if index_dtype == torch.uint32 and node_count > INT32_MAX + 1: + raise ValueError("canonical graph nodes require int64 indices") if graph.edge_vec.shape != (graph.source.shape[0], 3): - raise ValueError("canonical graph edge_vec must have shape (S, 3)") + raise ValueError("canonical edge_vec must have shape (S, 3) matching source") if graph.source.shape[0] < 2: raise ValueError("canonical graph edge storage must contain at least two slots") - if graph.source.shape[0] > UINT32_MAX: - raise ValueError("canonical graph edge storage exceeds the uint32 range") if graph.destination_row_ptr.shape != (node_count + 1,): raise ValueError("destination_row_ptr must have shape (N + 1,)") - if graph.source_row_ptr.shape != (node_count + 1,): - raise ValueError("source_row_ptr must have shape (N + 1,)") + if graph.source_row_ptr.ndim != 1: + raise ValueError("source_row_ptr must be one-dimensional") if graph.n_node.shape != graph.n_local.shape or graph.n_node.ndim != 1: raise ValueError("n_node and n_local must share shape (nf,)") - devices = { - graph.n_node.device, - graph.n_local.device, - graph.source.device, - graph.edge_vec.device, - graph.destination_row_ptr.device, - graph.source_row_ptr.device, - graph.source_order.device, - } - if len(devices) != 1: - raise ValueError("canonical graph tensors must reside on one device") tensors = ( graph.n_node, graph.n_local, @@ -112,6 +150,8 @@ def validate_canonical_graph_shapes( graph.source_row_ptr, graph.source_order, ) + if len({tensor.device for tensor in tensors}) != 1: + raise ValueError("canonical graph tensors must reside on one device") if not all(tensor.is_contiguous() for tensor in tensors): raise ValueError("canonical graph tensors must be contiguous") @@ -160,11 +200,9 @@ def canonical_graph_from_neighbor_graph( node_count = graph.destination_row_ptr.shape[0] - 1 storage_edge_count = max(physical_edge_count, 2) - if storage_edge_count > UINT32_MAX: - raise ValueError("canonical graph edge storage exceeds the uint32 range") source = torch.zeros( storage_edge_count, - dtype=torch.uint32, + dtype=canonical_index_dtype(node_count), device=graph.edge_index.device, ) edge_vec = torch.zeros( @@ -177,17 +215,17 @@ def canonical_graph_from_neighbor_graph( storage_edge_count, dtype=torch.int64, device=graph.edge_index.device, - ).to(torch.uint32) + ) if physical_edge_count: source[:physical_edge_count] = graph.edge_index[0, :physical_edge_count].to( - torch.uint32 + source.dtype ) edge_vec[:physical_edge_count] = graph.edge_vec[:physical_edge_count].to( torch.float32 ) source_order[:physical_edge_count] = graph.source_order[ :physical_edge_count - ].to(torch.uint32) + ].to(torch.int64) result = CanonicalGraph( n_node=graph.n_node.contiguous(), @@ -202,9 +240,75 @@ def canonical_graph_from_neighbor_graph( return result +def canonical_graph_from_pairs( + center: torch.Tensor, + source: torch.Tensor, + edge_vec: torch.Tensor, + n_node: torch.Tensor, + node_count: int, +) -> CanonicalGraph: + """Build a local canonical graph from active, unpadded neighbor pairs. + + Parameters + ---------- + center, source : torch.Tensor + Valid global node indices with shape (E,), int64, on one device. + edge_vec : torch.Tensor + Neighbor-minus-center vectors in angstrom with shape (E, 3), on the + index device. Geometry is converted to float32 for deployment. + n_node : torch.Tensor + Per-frame local node counts with shape (nf,), int64, on the same device. + node_count : int + Total number of local nodes across frames, equal to ``n_node.sum()``. + + Returns + ------- + CanonicalGraph + Destination-major graph with dual CSR and at least two storage slots. + All nodes are local; no coordinate differentiation is required. + """ + edge_count = source.shape[0] + destination_order = torch.argsort(center, stable=True) + center = center.index_select(0, destination_order) + source = source.index_select(0, destination_order) + edge_vec = edge_vec.index_select(0, destination_order).to(torch.float32) + boundaries = torch.arange(node_count + 1, dtype=torch.int64, device=source.device) + destination_row_ptr = torch.searchsorted(center, boundaries) + source_order = torch.argsort(source, stable=True) + source_row_ptr = torch.searchsorted( + source.index_select(0, source_order), boundaries + ) + # Guard slots keep the exported edge axis dynamic even for zero or one + # physical edge. CSR offsets exclude these slots from every node row. + if edge_count < 2: + padding = 2 - edge_count + source = torch.cat((source, source.new_zeros(padding))) + edge_vec = torch.cat((edge_vec, edge_vec.new_zeros((padding, 3)))) + source_order = torch.cat( + ( + source_order, + torch.arange(edge_count, 2, dtype=torch.int64, device=source.device), + ) + ) + n_node = n_node.contiguous() + return CanonicalGraph( + n_node=n_node, + n_local=n_node, + source=source.to(canonical_index_dtype(node_count)), + edge_vec=edge_vec, + destination_row_ptr=destination_row_ptr, + source_row_ptr=source_row_ptr, + source_order=source_order, + ) + + __all__ = [ - "UINT32_MAX", + "INT32_MAX", "CanonicalGraph", "canonical_graph_from_neighbor_graph", + "canonical_graph_from_pairs", + "canonical_index_bytes", + "canonical_index_dtype", + "canonical_index_words", "validate_canonical_graph_shapes", ] diff --git a/deepmd/pt_expt/utils/edge_schema.py b/deepmd/pt_expt/utils/edge_schema.py index b2cbad736d..9e8a3b2ab8 100644 --- a/deepmd/pt_expt/utils/edge_schema.py +++ b/deepmd/pt_expt/utils/edge_schema.py @@ -352,6 +352,8 @@ def merge_frame_edge_schemas(frames: list[EdgeNeighborList]) -> EdgeNeighborList """Merge per-frame local edge schemas into a batched flattened schema.""" if not frames: raise ValueError("at least one frame schema is required") + if len(frames) == 1: + return frames[0] coord = torch.cat([frame.coord for frame in frames], dim=0) atype = torch.cat([frame.atype for frame in frames], dim=0) nloc = coord.shape[1] diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index 3b117732a6..b19b1f6d35 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -34,18 +34,17 @@ def resolve_auto_graph_builder( ) -> str: """Resolve ``neighbor_graph_method="auto"`` to a concrete inference builder. - Single owner of the inference / DeepEval auto ladder. Training uses + Single owner of graph inference / DeepEval auto selection. Training uses :func:`resolve_neighbor_graph_method`, which never selects ``vesin``. - Mirrors :func:`deepmd.pt.model.model.sezm_model._select_neighbor_builder`: - ``vesin`` is eligible only for a single-frame batch (``nf == 1``), because - its API loops frames in Python (~1 ms/frame). Multi-frame batches stay on - ``nv`` (CUDA) or ``dense`` so ``auto_batch_size`` / ``dp test`` do not - regress to the per-frame loop. + ``vesin`` is preferred for single-frame CUDA inference. Its API loops over + frames in Python, so multi-frame CUDA batches prefer ``nv`` and fall back + to ``dense`` when Toolkit-Ops is unavailable. Notes ----- - * CUDA + ``nvalchemiops``: ``nv`` (any ``nf``). + * CUDA + ``nf == 1`` + ``vesin.torch``: ``vesin``. + * CUDA + ``nvalchemiops``: ``nv``. * CPU + the operator library: ``cell``, whose search is threaded. * ``nf == 1`` + ``vesin.torch``: ``vesin``. * otherwise: ``dense``. @@ -92,6 +91,8 @@ def resolve_auto_graph_builder( raise ValueError(f"nf must be >= 1, got {nf}") dev = torch.device(device) + if dev.type == "cuda" and nf == 1 and is_vesin_torch_available(): + return "vesin" nv_available = is_nv_available() if dev.type == "cuda" and nv_available: return "nv" diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index 9483b3ef66..23d35d9092 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -13,6 +13,8 @@ Any, ) +import numpy as np + from deepmd.dpmodel.utils.lmdb_data import ( DistributedLmdbBatchSampler, LmdbBatchIterator, @@ -150,6 +152,14 @@ def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: del sys_idx return next(self._batch_iterator) + def state_dict(self) -> dict[str, Any]: + """Return the next training-batch cursor, excluding statistical reads.""" + return self._batch_iterator.state_dict() + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore training progress without drawing or decoding any batch.""" + self._batch_iterator.load_state_dict(state) + def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: """Return one batch from a homogeneous statistical system. @@ -232,10 +242,41 @@ def get_nsystems(self) -> int: """Return one logical LMDB training dataset.""" return 1 + def get_model_sampling_size(self) -> int: + """Count retained frames, matching native PyTorch LMDB task weights.""" + return len(self._reader) + + @property + def uniform_sampling_mass(self) -> float: + """Return the queued batch's systems-per-frame selection weight.""" + return self._batch_iterator.uniform_sampling_mass + + @property + def uniform_nominal_mass(self) -> float: + """Return the pending pass's mean systems-per-frame selection weight.""" + return self._batch_iterator.uniform_nominal_mass + def get_batch_pass_length(self) -> int: """Return the local pass length belonging to the last returned batch.""" return self._batch_iterator.batch_pass_length + def get_epoch_batch_counts(self, world_size: int) -> np.ndarray: + """Return the actual local workload after distributed batch padding.""" + return np.asarray([len(self._sampler)], dtype=np.float64) + + def get_ntypes(self) -> int: + return len(self.type_map) + + def get_type_map(self) -> list[str]: + return self.type_map + + def get_batch_size(self) -> list[int]: + return [self._reader.batch_size] + + @property + def system_dirs(self) -> list[str]: + return [self._reader.lmdb_path] + @property def nbatches(self) -> list[int]: """Return the global batch count of one full pass.""" diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py index aacd1383bc..663cff1aaf 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -359,5 +359,12 @@ def build_neighbor_graph_nv( if pair_excl is not None: graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) if with_csr or canonicalize: - graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize) + # Matrix slots are decoded in ascending center order; compact filtering + # preserves that order, whereas an exclusion mask can introduce holes. + graph = attach_edge_csr( + graph, + nf * nloc, + canonicalize=canonicalize, + destination_sorted=pair_excl is None or compact, + ) return graph diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index b3edf999ad..e5fb598494 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import contextlib import ctypes +import dataclasses import json import logging import os @@ -14,6 +15,10 @@ import numpy as np import torch +from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, + flatten_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -27,6 +32,7 @@ traverse_model_dict, ) from deepmd.pt.utils.compile_compat import ( + _torch_release, traced_output_keys, ) from deepmd.pt_expt.model.graph_lower import ( @@ -145,6 +151,32 @@ def _graph_reads_source_csr(exported: Any) -> bool: return not found +def _canonical_requires_source_csr(model: Any, lower_kind: str) -> bool: + """Return whether a compact canonical lower reads the source CSR on every call. + + The force assembly of compressed DPA1 walks the source CSR, and so does + the composition of operators that evaluates a long-range module whose + charge network the fused long-range operator does not reproduce. The + fused DPA4C operators assemble the forces in fixed point and read the + source CSR only for the per-atom virial, so a caller that declines it may + pass two-entry placeholders in its place. + """ + if lower_kind == "dpa1_canonical": + return True + decomposition = model.atomic_model.fused_decomposition() + learned = None if decomposition is None else decomposition[0] + if learned is None: + # The compact lower serves no such model, which the export reports. + return True + if not learned.has_long_range(): + return False + from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + fused_long_range_eligible, + ) + + return not fused_long_range_eligible(learned.long_range) + + def _numpy_to_json_serializable(model_obj: dict) -> dict: """Convert numpy arrays in a model dict to JSON-serializable lists.""" return traverse_model_dict( @@ -307,6 +339,91 @@ def check_graph_trace_torch_version(model: torch.nn.Module) -> None: ) +def mesh_freeze_supported() -> bool: + """Whether AOTInductor runs the mesh evaluation of a periodic graph lower. + + The graph lower of a periodic long-range model evaluates its reciprocal + sum by FFTs over the mesh of every cell, and the derivatives of the real + transforms allocate complex buffers through the proxy executor of + AOTInductor. Its runtime allocates them with their declared dtype from + PyTorch 2.12 on; older runtimes allocate booleans, which the + complex-to-complex transform of the derivative rejects when the packaged + model runs. + """ + return _torch_release() >= (2, 12) + + +def mesh_compile_supported(device: torch.device) -> bool: + """Whether ``torch.compile`` forms the virial of the mesh evaluation on ``device``. + + The CPU code generation of Inductor forms the strain derivative of the + mesh energy in a compiled training graph correctly from PyTorch 2.13 on, + the CUDA code generation in every release that compiles the model. + """ + return device.type != "cpu" or _torch_release() >= (2, 13) + + +def check_mesh_freeze_torch_version(metadata: dict) -> None: + """Fail fast before AOTInductor packages a mesh evaluation it cannot run. + + Parameters + ---------- + metadata : dict + The metadata of the traced lower, with its ``lower_input_kind`` and + its ``topologies``. + + Raises + ------ + RuntimeError + If the graph lower reads the PME mesh and :func:`mesh_freeze_supported` + does not hold. + """ + if ( + metadata.get("lower_input_kind") == "graph" + and PME_MESH in metadata.get("topologies", ()) + and not mesh_freeze_supported() + ): + raise RuntimeError( + "freezing a periodic long-range model into a .pt2 graph lower " + f"requires torch >= 2.12; found torch {torch.__version__}, whose " + "AOTInductor runtime allocates the complex buffers of the mesh " + "derivatives as booleans. Upgrade torch, freeze to a .pte file, " + "or compress the model for the long-range operator on a CUDA " + "device." + ) + + +def check_mesh_compile_torch_version( + model: torch.nn.Module, device: torch.device +) -> None: + """Fail fast before compiling a model whose virial Inductor forms wrongly. + + Parameters + ---------- + model : torch.nn.Module + The model about to be traced for compiled training. + device : torch.device + The device of the trace. + + Raises + ------ + RuntimeError + If the model reads the PME mesh and :func:`mesh_compile_supported` + does not hold on ``device``. + """ + if ( + PME_MESH in model.atomic_model.required_topologies() + and not mesh_compile_supported(device) + ): + raise RuntimeError( + "compiled training of a periodic long-range model on the PME mesh " + f"requires torch >= 2.13 on the CPU; found torch {torch.__version__}, " + "whose Inductor CPU code generation forms a wrong virial for it. " + "Upgrade torch, train on a CUDA device, or set " + "`training.enable_compile` to false." + ) + + # Module-level cache for the trace-time sendlist buffer. The pointer # value embedded in ``send_list_tensor`` references this numpy array's # data; the array must outlive the trace + export call. Caching here @@ -501,7 +618,7 @@ def build_synthetic_graph_inputs( Single source of the trace-time graph inputs, shared by ``.pt2`` export (:func:`_trace_and_export`) and compiled training (:func:`deepmd.pt_expt.train.training._trace_and_compile_graph`). Builds a - small random system and runs the carry-all + small deterministic system and runs the carry-all :func:`~deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph` with a padded ``GraphLayout(edge_capacity=e_max)`` trace sample. Export uses the destination-major graph with CSR metadata; compiled training selects the @@ -519,8 +636,10 @@ def build_synthetic_graph_inputs( (before ``fparam``/``aparam``); there is no ``charge_spin`` slot at all (native spin rejects ``add_chg_spin_ebd`` at build). - The system (``rng(42)``, ``box = rcut*3``, centered coords, ``atype[:, i] = - i % ntypes``) is identical for both callers. + The system is identical for both callers and independent of dense neighbor + capacities. For at least four atoms it includes an isolated center and a + connected three-atom cluster. The highest atom type occurs in the cluster + even when the type map is larger than the sample. Parameters ---------- @@ -587,9 +706,21 @@ def build_synthetic_graph_inputs( box_np = (np.eye(3, dtype=np.float64) * box_size).reshape(1, 9) rng = np.random.default_rng(42) coord_np = rng.random((nframes, nloc, 3)) * box_size * 0.5 + box_size * 0.25 + if nloc >= 4: + # Exercise multi-edge aggregation and empty neighborhoods in the same + # trace. Every other atom lies in [.75*rcut, 2.25*rcut]^3, so the + # origin is isolated even under periodic minimum-image distances. + coord_np[:, [0, 2, 3]] = rcut * np.array( + [[1.5, 1.5, 1.5], [1.6, 1.5, 1.5], [1.5, 1.65, 1.5]] + ) + # Keep the isolated center inside n_local (the final atom is a + # synthetic ghost), including the minimal four-atom fixture. + coord_np[:, 1] = 0.0 atype_np = np.zeros((nframes, nloc), dtype=np.int64) for i in range(nloc): atype_np[:, i] = i % ntypes + # Include the upper type-index boundary without growing the node axis. + atype_np[:, 0] = ntypes - 1 coord_t = torch.tensor(coord_np, dtype=dtype, device=device) atype_t = torch.tensor(atype_np, dtype=torch.int64, device=device) @@ -674,24 +805,126 @@ def build_synthetic_graph_inputs( ) +def build_synthetic_topologies( + model: torch.nn.Module, + n_node_total: int, + n_entry: int, + *, + dtype: torch.dtype, + device: torch.device, + nframes: int, + forbidden: set[int] | frozenset[int] = frozenset(), +) -> dict[str, Any]: + """Build trace-time samples of the topologies a model declares. + + Single source of the trace-time topology inputs, shared by graph export + and compiled training. Tracing reads only their shapes. The pair + topology lists ``n_entry`` pairs, a length the caller keeps distinct from + the other dynamic axes; a mesh takes the three smallest even dimensions, + and a k grid the three smallest lengths for its slots, its blocks and the + points of a block, of at least four that no other axis and no entry of + ``forbidden`` has. Every length thus differs from the two rows of a pair + index and the three components of a vector, and from every other dynamic + axis, which a compiled graph would otherwise identify with it. + + Parameters + ---------- + model : torch.nn.Module + The model; its atomic model declares the topologies. + n_node_total : int + Length of the node axis. + n_entry : int + Length of the pair axis. + dtype : torch.dtype + Precision of the geometry. + device : torch.device + Target device. + nframes : int + Length of the frame axis. + forbidden : set[int] + Sizes of the static dimensions of the model, which the mesh and k-grid + axes avoid. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + A sample of every declared topology; empty when none is declared. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + K_GRID, + KGrid, + PairGraph, + ) + from deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph import ( + mesh_grid_of_dims, + ) + + taken = set(forbidden) | {nframes, n_node_total, n_node_total + 1, n_entry} + mesh_dims = [] + for size in range(4, 4 + 2 * len(taken) + 6, 2): + if size not in taken and len(mesh_dims) < 3: + mesh_dims.append(size) + free = sorted(set(range(4, len(taken) + 10)) - taken - set(mesh_dims)) + n_slot, n_block, n_point = free[:3] + coord = torch.ones(n_node_total, 3, dtype=dtype, device=device) + # Contiguous like the inputs at run time: a compiled graph reads the + # strides of its samples. + box = torch.eye(3, dtype=dtype, device=device).repeat(nframes, 1, 1) + topologies = {} + for name in model.atomic_model.required_topologies(): + if name == K_GRID: + slot = torch.arange(nframes * n_slot, dtype=torch.int64, device=device) + blocks = (nframes, n_block, n_point) + topologies[name] = KGrid( + # Nonzero wave numbers: the kernel transform carries 1/k^2. + vec=torch.ones(*blocks, 3, dtype=dtype, device=device), + mask=torch.ones(*blocks, dtype=torch.bool, device=device), + slot=(slot % n_node_total).reshape(nframes, n_slot), + volume=torch.ones(nframes, dtype=dtype, device=device), + coord=coord, + strain=torch.zeros(nframes, 3, 3, dtype=dtype, device=device), + ) + elif name == PME_MESH: + topologies[name] = mesh_grid_of_dims( + box, + coord, + tuple(mesh_dims), + torch.ones(nframes, dtype=torch.bool, device=device), + ) + else: + index = torch.arange(2 * n_entry, dtype=torch.int64, device=device) + topologies[name] = PairGraph( + index=index.reshape(2, n_entry) % n_node_total, + vec=torch.ones(n_entry, 3, dtype=dtype, device=device), + mask=torch.ones(n_entry, dtype=torch.bool, device=device), + ) + return topologies + + def build_synthetic_canonical_graph_inputs( model: torch.nn.Module, e_max: int, *, device: torch.device, want_spin: bool = False, + want_long_range: bool = False, ) -> tuple[torch.Tensor, ...]: """Build compact canonical trace inputs for compressed CUDA descriptors. ``want_spin`` appends the per-node moment at slot 8, the last slot of the compact ABI, matching :meth:`~deepmd.pt_expt.model.native_spin_model.NativeSpinEnergyModel.forward_lower_canonical_graph_exportable`. + ``want_long_range`` appends the inputs that the compact lower of a model + with a long-range module reads: the float64 node positions at slot 8, the + int64 communicator, a CPU tensor of one element, at slot 9, and the + float64 cells of the frames, a CPU tensor of shape ``(nf, 9)``, at slot 10. """ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) from deepmd.pt_expt.utils.canonical_graph import ( canonical_graph_from_neighbor_graph, + canonical_index_bytes, ) sample = build_synthetic_graph_inputs( @@ -736,45 +969,59 @@ def build_synthetic_canonical_graph_inputs( atype, compact.n_node, compact.n_local, - compact.source, + canonical_index_bytes(compact.source), compact.edge_vec, compact.destination_row_ptr, compact.source_row_ptr, compact.source_order, ) - return compact_sample if spin is None else (*compact_sample, spin) + if spin is not None: + return (*compact_sample, spin) + if want_long_range: + # Tracing reads the shape of the positions alone; distinct values keep + # the sample nodes apart. The communicator of a single domain is zero. + coord = 0.7 * torch.arange( + 3 * atype.shape[0], dtype=torch.float64, device=device + ).reshape(-1, 3) + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") + box = torch.zeros(n_node.shape[0], 9, dtype=torch.float64, device="cpu") + return (*compact_sample, coord, communicator, box) + return compact_sample def _build_canonical_graph_dynamic_shapes( *sample_inputs: torch.Tensor, ) -> tuple: - """Build dynamic shapes for the compact deployment ABI. + """Build dynamic node, edge and index-word axes for the compact ABI. - The trailing spin slot is present only for a native-spin model, so the - sample length selects the shape tuple. + The trailing slots, the moments of a native-spin model or the positions, + the communicator and the cells of a model with a long-range module, are + present only for such models, so the sample length selects the shape + tuple. The communicator holds one address and is static; the cells follow + the frame axis. """ - from deepmd.pt_expt.utils.canonical_graph import ( - UINT32_MAX, - ) - nframes_dim = torch.export.Dim("nframes", min=1) node_dim = torch.export.Dim("n_node_total", min=1) - edge_storage_dim = torch.export.Dim( - "nedge_storage", - min=2, - max=UINT32_MAX, - ) + edge_storage_dim = torch.export.Dim("nedge_storage", min=2) + # The source CSR keeps extents of its own: a caller that declines the + # per-atom virial passes two-entry placeholders in its place. + source_row_dim = torch.export.Dim("n_source_row", min=2) + source_order_dim = torch.export.Dim("n_source_order", min=2) + word_bytes_dim = torch.export.Dim("index_word_bytes", min=4, max=8) shapes = ( {0: node_dim}, {0: nframes_dim}, {0: nframes_dim}, + {0: edge_storage_dim, 1: word_bytes_dim}, {0: edge_storage_dim}, - {0: edge_storage_dim}, - {0: node_dim + 1}, {0: node_dim + 1}, - {0: edge_storage_dim}, + {0: source_row_dim}, + {0: source_order_dim}, ) - return shapes if len(sample_inputs) == len(shapes) else (*shapes, {0: node_dim}) + trailing = ({0: node_dim}, None, {0: nframes_dim})[ + : len(sample_inputs) - len(shapes) + ] + return (*shapes, *trailing) def count_synthetic_graph_edges( @@ -827,9 +1074,53 @@ def count_synthetic_graph_edges( return int(edge_mask.sum()) +def _topology_dynamic_shapes( + topology_names: tuple[str, ...], shared: dict[str, Any] +) -> list[dict[int, Any]]: + """Dynamic-shape specifications of the fields of the topologies of a graph lower. + + Every field follows the ``AXES`` its dataclass declares. The frame and + node axes are those of the other inputs in ``shared``; the entry axis of + every topology is a dimension of its own, of at least the two guard + entries, and so are the slot axis and the point axis of a k grid; a mesh + axis is twice a dimension, since the mesh dimensions are even. + + Parameters + ---------- + topology_names : tuple[str, ...] + The topologies, in the order of their inputs. + shared : dict[str, Any] + The dimensions of the ``"frame"`` and ``"node"`` axes. + + Returns + ------- + list[dict[int, Any]] + One specification per field, in the layout of flatten_topologies. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + TOPOLOGY_TYPES, + ) + + specs = [] + for name in topology_names: + cls = TOPOLOGY_TYPES[name] + dims = dict(shared) + dims["entry"] = torch.export.Dim(f"{name}_entries", min=2) + dims["slot"] = torch.export.Dim(f"{name}_slots", min=2) + dims["k_point"] = torch.export.Dim(f"{name}_block_points", min=2) + for axis in ("mesh_1", "mesh_2", "mesh_3"): + dims[axis] = 2 * torch.export.Dim(f"{name}_{axis}_half", min=1) + specs += [ + {dim: dims[axis] for dim, axis in cls.AXES[field.name].items()} + for field in dataclasses.fields(cls) + ] + return specs + + def _build_graph_dynamic_shapes( *sample_inputs: torch.Tensor | None, is_native_spin: bool = False, + topology_names: tuple[str, ...] = (), ) -> tuple: """Build dynamic-shape specifications for the graph-form forward_lower export. @@ -857,6 +1148,9 @@ def _build_graph_dynamic_shapes( Whether ``sample_inputs`` follows the native-spin positional ABI (spin at slot 10) instead of the regular energy ABI (charge_spin at slot 12). + topology_names : tuple[str, ...] + Topologies whose fields follow the inputs above (see + :func:`_topology_dynamic_shapes`). """ nframes_dim = torch.export.Dim("nframes", min=1) n_node_total_dim = torch.export.Dim("n_node_total", min=1) @@ -873,6 +1167,9 @@ def _build_graph_dynamic_shapes( {0: nedge_dim}, # source_order: (E,) {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) ) + topologies = _topology_dynamic_shapes( + topology_names, {"frame": nframes_dim, "node": n_node_total_dim} + ) if is_native_spin: spin = sample_inputs[10] fparam = sample_inputs[11] @@ -886,6 +1183,7 @@ def _build_graph_dynamic_shapes( {0: nframes_dim} if fparam is not None else None, # fparam {0: n_node_total_dim} if aparam is not None else None, # aparam {0: nframes_dim} if charge_spin is not None else None, # charge_spin + *topologies, ) fparam = sample_inputs[10] aparam = sample_inputs[11] @@ -898,6 +1196,7 @@ def _build_graph_dynamic_shapes( # dim would make torch.export prove/reject the equality). {0: n_node_total_dim} if aparam is not None else None, # aparam {0: nframes_dim} if charge_spin is not None else None, # charge_spin + *topologies, ) @@ -1198,7 +1497,7 @@ def _collect_metadata( } if is_spin: # The scheme is what selects the serving backend class in C++ - # (``deepmd_create_deepspin_backend_v1``): "native" is served by + # (``deepmd_create_deepspin_backend_v2``): "native" is served by # NativeSpinPTExpt, "deepspin" by DeepSpinPTExpt. It is a property of # the model alone, orthogonal to ``lower_input_kind`` below. meta["spin_scheme"] = spin_scheme @@ -1257,7 +1556,11 @@ def _probe_has_message_passing(obj: object) -> bool | None: meta["lower_input_kind"] = lower_kind meta["graph_edge_dtype"] = graph_edge_dtype(model, lower_kind) if lower_kind in ("dpa1_canonical", "dpa4c_canonical"): - meta["canonical_index_dtype"] = "uint32" + meta["canonical_index_dtype"] = "auto" + meta["canonical_index_storage"] = "byte_words" + meta["canonical_requires_source_csr"] = _canonical_requires_source_csr( + model, lower_kind + ) # Model-level pair-type exclusion (``pair_exclude_types``): a list of # ``[ti, tj]`` type pairs whose interaction is dropped. Exclusion is a @@ -1286,6 +1589,16 @@ def _probe_has_message_passing(obj: object) -> bool | None: pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] break meta["pair_exclude_types"] = pair_exclude_types + # Topologies the lower reads beside the neighbor graph, and the parameters + # they are built with; absent when the model reads none. The graph lower + # takes the fields of each after its other inputs (see + # ``flatten_topologies``), built by the caller from the coordinates and + # the cell of each frame. The compact canonical lower takes the node + # positions instead, from which its operators evaluate the long range. + topologies = model.atomic_model.required_topologies() + if topologies: + meta["topologies"] = list(topologies) + meta["topology_params"] = model.atomic_model.topology_params() return meta @@ -1305,7 +1618,8 @@ def serialize_from_file(model_file: str) -> dict: The serialized model data. If the archive contains ``model_def_script.json`` (training config), it is included under the ``"model_def_script"`` key. ``lower_input_kind`` records - the concrete lower ABI from the artifact metadata. + the concrete lower ABI and ``do_atomic_virial`` whether the archive + exports the per-atom virial, both from the artifact metadata. """ if model_file.endswith(".pt2"): return _serialize_from_file_pt2(model_file) @@ -1320,7 +1634,10 @@ def _serialize_from_file_pte(model_file: str) -> dict: "model_def_script.json": "", "metadata.json": "", } - torch.export.load(model_file, extra_files=extra_files) + # Storage reconstruction requires CPU staging independently of the caller's + # default device; tensor placement is restored from the archive metadata. + with torch.device("cpu"): + torch.export.load(model_file, extra_files=extra_files) model_dict = json.loads(extra_files["model.json"]) model_dict = _json_to_numpy(model_dict) metadata = ( @@ -1329,6 +1646,7 @@ def _serialize_from_file_pte(model_file: str) -> dict: model_dict["lower_input_kind"] = metadata.get( "lower_input_kind", model_dict.get("lower_input_kind", "nlist") ) + model_dict["do_atomic_virial"] = bool(metadata.get("do_atomic_virial", False)) if extra_files["model_def_script.json"]: model_dict["model_def_script"] = json.loads( extra_files["model_def_script.json"] @@ -1366,6 +1684,7 @@ def _serialize_from_file_pt2(model_file: str) -> dict: model_dict["lower_input_kind"] = metadata.get( "lower_input_kind", model_dict.get("lower_input_kind", "nlist") ) + model_dict["do_atomic_virial"] = bool(metadata.get("do_atomic_virial", False)) if model_def_script_json: model_dict["model_def_script"] = json.loads(model_def_script_json) return model_dict @@ -1545,6 +1864,32 @@ def _dpa4_kernel_levels_for_target( os.environ[name] = value +def _fold_case_embeddings(model: torch.nn.Module) -> None: + """Fold the fixed case embedding of every fitting into its first layer. + + A deployed model evaluates a single case, so the schema selection, the + launch tuners and the fused fitting operators see a plain fitting network + (see :meth:`~deepmd.dpmodel.fitting.general_fitting.GeneralFitting.fold_case_embedding`). + """ + from deepmd.dpmodel.fitting.general_fitting import ( + GeneralFitting, + ) + + # The fold replaces submodules, so the fittings are collected before it. + fittings = [ + module for module in model.modules() if isinstance(module, GeneralFitting) + ] + for fitting in fittings: + fitting.fold_case_embedding() + if fitting.dim_case_embd: + log.warning( + "The case embedding of a fitting network is kept, since the " + "residual of its first layer adds the input to the output; " + "the export evaluates that network without the fused fitting " + "operators." + ) + + def _select_graph_lower_kind(data: dict, *, allow_canonical: bool) -> str | None: """Select the graph schema supported by the target model. @@ -1569,6 +1914,7 @@ def _select_graph_lower_kind(data: dict, *, allow_canonical: bool) -> str | None ) model = BaseModel.deserialize(data["model"]) + _fold_case_embeddings(model) if not (model_uses_graph_lower(model) and _supports_graph_export(model)): return None if allow_canonical: @@ -1647,6 +1993,8 @@ def deserialize_to_file( model_json_override: dict | None = None, do_atomic_virial: bool = False, lower_kind: str = "nlist", + *, + with_comm: bool = True, ) -> None: """Deserialize a dictionary to a .pte or .pt2 model file. @@ -1690,6 +2038,10 @@ def deserialize_to_file( explicitly selects other levels. The selected schema is recorded as ``lower_input_kind`` in ``metadata.json``. + with_comm : bool, optional + Include a communication artifact in ``.pt2`` exports for models with + cross-rank message passing, by default True. Disable for + single-process inference. """ lower_kind = _resolve_target_lower_kind(model_file, data, lower_kind) uses_dpa4_defaults = model_file.endswith(".pt2") and _uses_dpa4_kernel_defaults( @@ -1717,6 +2069,7 @@ def deserialize_to_file( model_json_override, do_atomic_virial, lower_kind, + with_comm=with_comm, ) else: _deserialize_to_file_pte( @@ -1728,6 +2081,20 @@ def deserialize_to_file( ) +def _canonical_symbolic_shapes(lower_kind: str) -> contextlib.AbstractContextManager: + """Preserve zero and one in the canonical program's dynamic shape ranges. + + A single frame is a valid deployment input, including graphs with zero + or one physical edge. Size-oblivious symbols retain these extents through + tracing and AOTInductor compilation without disabling runtime guards. + """ + if lower_kind not in ("dpa1_canonical", "dpa4c_canonical"): + return contextlib.nullcontext() + from torch.fx.experimental import _config as symbolic_config + + return symbolic_config.patch(backed_size_oblivious=True) + + def _trace_and_export( data: dict, model_json_override: dict | None = None, @@ -1739,7 +2106,10 @@ def _trace_and_export( import deepmd.pt_expt.utils.env as _env target_device = _env.DEVICE - with _dpa4_kernel_levels_for_target(data["model"], target_device): + with ( + _dpa4_kernel_levels_for_target(data["model"], target_device), + _canonical_symbolic_shapes(lower_kind), + ): return _trace_and_export_impl( data, model_json_override, @@ -1750,6 +2120,70 @@ def _trace_and_export( ) +@contextlib.contextmanager +def _cpu_reference_descriptors(model: torch.nn.Module) -> Iterator[None]: + """Temporarily use CPU reference descriptors for the eager vacuum fold. + + Deployment descriptors bind accelerator operators at construction, before + they move to CPU for fake tracing. Changing their device or the kernel + environment cannot unbind those operators. Rebuild only descriptors whose + fittings still need a reference, leaving the folded bias and non-persistent + vacuum table on the original fitting. Restore the deployment descriptors + before tracing so their optimized bindings remain intact. + """ + from deepmd.dpmodel.atomic_model.dp_atomic_model import ( + DPAtomicModel, + ) + from deepmd.pt_expt.utils import ( + env, + ) + + atomic_models = [ + module + for module in model.modules() + if isinstance(module, DPAtomicModel) + and module.fitting_net.needs_vacuum_descriptor() + ] + if not atomic_models: + yield + return + + levels = ( + "DP_CUDA_INFER", + "DP_TRITON_INFER", + "DP_CUTILE_INFER", + "DP_CUTE_INFER", + "DP_CUDA_TRAIN", + "DP_TRITON_TRAIN", + "DP_AMP_INFER", + ) + saved_levels = {name: os.environ.get(name) for name in levels} + saved_device = env.DEVICE + descriptors = [] + try: + env.DEVICE = torch.device("cpu") + os.environ.update(dict.fromkeys(levels, "0")) + for atomic_model in atomic_models: + descriptor = atomic_model.descriptor + reference = type(descriptor).deserialize(descriptor.serialize()).eval() + # Virtual-spin backbones can disable this route at runtime; that + # choice is deliberately absent from the portable descriptor. + if not descriptor.uses_graph_lower() and reference.uses_graph_lower(): + reference.disable_graph_lower() + descriptors.append((atomic_model, descriptor)) + atomic_model.descriptor = reference + yield + finally: + for atomic_model, descriptor in descriptors: + atomic_model.descriptor = descriptor + env.DEVICE = saved_device + for name, value in saved_levels.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + def _trace_and_export_impl( data: dict, model_json_override: dict | None = None, @@ -1832,13 +2266,28 @@ def _trace_and_export_impl( # Registry-dispatched (incl. native spin, type "native_spin"): the # pt_expt BaseModel registry returns this backend's torch class. model = BaseModel.deserialize(data["model"]) + # The fold rebuilds the first fitting layers on the default device, where + # the deserialized model resides, so it precedes the move below. + _fold_case_embeddings(model) model.eval() - # The vacuum reference is resolved on the export object, folded into the - # fitting bias or stored as a per-type table, so the exported graph - # carries no reference atoms; the move to the tracing device follows, so - # the table lands there with the rest of the model. - model.fold_vacuum_reference() + # A frozen model evaluates the reciprocal sum of a periodic frame on the + # mesh of its cell, whatever the mode it was trained in. + from deepmd.dpmodel.long_range.latent_charge import ( + LatentChargeLongRange, + ) + + for module in model.modules(): + if isinstance(module, LatentChargeLongRange): + module.use_mesh_evaluation() + # Fold the vacuum reference on the CPU tracing device. vacuum_descriptor + # evaluates an isolated atom of every type through the descriptor; on CUDA + # that path can allocate several GiB even when export already holds most of + # the GPU. Move first so the fold runs where tracing will, and so the + # folded bias or vacuum table stays on that device with the rest of the + # model. model.to("cpu") + with torch.no_grad(), _cpu_reference_descriptors(model): + model.fold_vacuum_reference() # Device-dependent Python branches resolve on the CPU tracing inputs, so # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2) @@ -1870,7 +2319,7 @@ def _trace_and_export_impl( run_autotune, ) - run_autotune(model, target_device) + run_autotune(model, target_device, lower_kind) # Pack after checkpoint deserialization. The non-persistent buffers become # constants in the CPU trace and move with the exported graph to the @@ -1999,11 +2448,17 @@ def _trace_and_export_impl( while e_sample in (_forbidden | {nframes_sample, n_sample}): e_sample += 1 if canonical: + # The compact lower of a model with a long-range module reads the + # node positions and the cells, from which its operator evaluates + # the long range, and the communicator of the domains of a + # decomposed frame. + want_long_range = bool(model.atomic_model.required_topologies()) sample_inputs = build_synthetic_canonical_graph_inputs( model, e_sample, device=torch.device("cpu"), want_spin=is_native_spin, + want_long_range=want_long_range, ) traced = model.forward_lower_canonical_graph_exportable( *sample_inputs, @@ -2085,6 +2540,20 @@ def _trace_and_export_impl( device=torch.device("cpu"), want_spin=is_native_spin, ) + # The pair axis is dynamic; its sample length is kept distinct + # from the other axes and from the CSR offsets. + topologies = build_synthetic_topologies( + model, + n_sample, + next_safe_prime( + e_sample + 2, + _forbidden | {nframes_sample, n_sample, n_sample + 1, e_sample}, + ), + dtype=edge_dtype, + device=torch.device("cpu"), + nframes=nframes_sample, + forbidden=_forbidden | {e_sample}, + ) if is_native_spin: # Native-spin ABI (NativeSpinEnergyModel.forward_lower_graph_exportable): # slot 10 is ``spin`` (mandatory), slots 11/12 are @@ -2110,12 +2579,16 @@ def _trace_and_export_impl( do_atomic_virial=do_atomic_virial, charge_spin=sample_inputs[12], destination_sorted=True, + topologies=topologies, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) dynamic_shapes = _build_graph_dynamic_shapes( - *sample_inputs, is_native_spin=is_native_spin + *sample_inputs, + is_native_spin=is_native_spin, + topology_names=tuple(topologies), ) + sample_inputs = (*sample_inputs, *flatten_topologies(topologies)) output_keys = traced_output_keys(traced) exported = torch.export.export( traced, @@ -2446,7 +2919,81 @@ def _charge_state_descriptor(data: dict, metadata: dict) -> Any | None: return descriptor -def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: +def _charge_state_long_range(data: dict, metadata: dict) -> Any | None: + """Return the long-range module folded to the charge state, if any. + + Compressing the atomic model folds a long-range charge constraint that + reads its total charge from the charge state to the condition of the + descriptor, so the fold rebuilds the state of the module with the + artifacts of the descriptor. Only a model with a long-range module, which + the collected metadata names by its topologies, pays for rebuilding the + model. + + Parameters + ---------- + data + Serialized model dictionary, as passed to the export entry point. + metadata + Archive metadata collected from that same model. + + Returns + ------- + Any or None + The evaluated long-range module, or ``None`` when the model folds no + long-range state. + """ + if not metadata.get("topologies"): + return None + + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + model = BaseModel.deserialize(data["model"]) + parts = model.atomic_model.fused_decomposition() + long_range = None if parts is None else parts[0].long_range + if long_range is None or not long_range.charge_state_constants(): + return None + model.to("cpu") + model.eval() + return long_range + + +class _ChargeStateRebuild(torch.nn.Module): + """Rebuild the charge state of a descriptor and of a long-range module. + + The outputs are those of the descriptor fold followed by the folded state + of the long-range module. + + Parameters + ---------- + descriptor + Compressed charge-conditioned descriptor, as returned by + :func:`_charge_state_descriptor`. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`. + """ + + def __init__(self, descriptor: Any, long_range: Any) -> None: + super().__init__() + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + ChargeStateFold, + ) + + self.descriptor = ChargeStateFold(descriptor) + self.long_range = long_range + + def forward(self, charge_spin: torch.Tensor) -> tuple[torch.Tensor, ...]: + return ( + *self.descriptor(charge_spin), + *self.long_range.rebuild_charge_state(charge_spin), + ) + + +def _match_charge_state_constants( + descriptor: Any, exported: Any, long_range: Any | None = None +) -> list[str]: """Name, per fold output, the constant of ``exported`` it replaces. A compressed charge-conditioned descriptor carries its frame condition in @@ -2479,6 +3026,10 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: :func:`_charge_state_descriptor`. exported Exported lower whose constants are to be named. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`, whose state is matched after the + artifacts of the descriptor; ``None`` when there is none. Returns ------- @@ -2499,9 +3050,18 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: for name, tensor in exported.state_dict.items() if isinstance(tensor, torch.Tensor) } + artifacts = [ + (name, getattr(descriptor, f"compress_{name}")) + for name in CHARGE_STATE_ARTIFACTS + ] + if long_range is not None: + artifacts += [ + ("long-range charge state", state) + for state in long_range.charge_state_constants() + ] constants: list[str] = [] - for name in CHARGE_STATE_ARTIFACTS: - artifact = getattr(descriptor, f"compress_{name}").detach().cpu() + for name, tensor in artifacts: + artifact = tensor.detach().cpu() if artifact.numel() == 0: constants.append("") continue @@ -2526,6 +3086,7 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: def _compile_charge_state_fold( descriptor: Any, aoti_configs: dict, + long_range: Any | None = None, ) -> bytes: """Compile the rebuild of the charge-state artifacts as its own archive. @@ -2541,6 +3102,10 @@ def _compile_charge_state_fold( :func:`_charge_state_descriptor`. aoti_configs Inductor options the inference lower was compiled with. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`, whose state the fold rebuilds after + the artifacts of the descriptor; ``None`` when there is none. Returns ------- @@ -2567,7 +3132,12 @@ def _compile_charge_state_fold( dtype=torch.float32, device="cpu", ) - fold = torch.export.export(ChargeStateFold(descriptor), (sample,)) + module = ( + ChargeStateFold(descriptor) + if long_range is None + else _ChargeStateRebuild(descriptor, long_range) + ) + fold = torch.export.export(module, (sample,)) if _env.DEVICE.type != "cpu": from torch.export.passes import ( move_to_device_pass, @@ -2587,6 +3157,8 @@ def _deserialize_to_file_pt2( model_json_override: dict | None = None, do_atomic_virial: bool = False, lower_kind: str = "nlist", + *, + with_comm: bool = True, ) -> None: """Deserialize a dictionary to a .pt2 model file (AOTInductor). @@ -2597,11 +3169,12 @@ def _deserialize_to_file_pt2( For models whose descriptor reports ``has_message_passing_across_ranks() == True`` (DPA2, DPA3 with ``use_loc_mapping=False``, or hybrids wrapping such children), - compiles a SECOND ``with-comm`` artifact and packs it alongside the - regular one. The ``with-comm`` variant accepts comm-dict tensors as - additional positional inputs and drives MPI ghost-atom exchange via - ``deepmd_export::border_op``. The C++ ``DeepPotPTExpt`` loader picks - the artifact based on the LAMMPS rank count at runtime. + compiles a SECOND ``with-comm`` artifact when ``with_comm`` is True and + packs it alongside the regular one. The ``with-comm`` variant accepts + comm-dict tensors as additional positional inputs and drives MPI + ghost-atom exchange via ``deepmd_export::border_op``. The C++ + ``DeepPotPTExpt`` loader picks the artifact based on the LAMMPS rank + count at runtime. Layout inside the .pt2 ZIP (PyTorch 2.11 strict layout): regular → artifact at ``model/`` (AOTInductor's own layout) @@ -2630,6 +3203,10 @@ def _deserialize_to_file_pt2( lower_kind=lower_kind, ) metadata["output_keys"] = output_keys + metadata["has_comm_artifact"] = with_comm and bool( + metadata.get("has_comm_artifact") + ) + check_mesh_freeze_torch_version(metadata) # A charge-state update swaps constants in one compiled artifact. A model # that also needs the with-comm lower would own two independent constant @@ -2686,25 +3263,39 @@ def _deserialize_to_file_pt2( aoti_configs["assert_indirect_indexing"] = False if is_cuda: aoti_configs["realize_opcount_threshold"] = 0 + # In the compact canonical lower of a compressed descriptor, expressions of + # the weights alone, such as the per-type energy bias that the program sums + # from two parameters ahead of the fused operator, are evaluated once, by + # the first call after the constants are loaded, instead of at every call. + # An update of the constants marks them for evaluation again at the next + # call, so a charge state applied at run time stays consistent. The setting + # applies to the compact canonical lower alone, the program whose constant + # updates the tests check; the other packages of the archive keep the + # shared options. + lower_configs = dict(aoti_configs) + if metadata.get("lower_input_kind") in ("dpa1_canonical", "dpa4c_canonical"): + lower_configs["aot_inductor.use_runtime_constant_folding"] = True log.info( "Compiling the AOTInductor package for %s (the slowest freeze stage; " "typically several minutes)...", _env.DEVICE, ) - aoti_compile_and_package( - exported, package_path=model_file, inductor_configs=aoti_configs - ) + with _canonical_symbolic_shapes(lower_kind): + aoti_compile_and_package( + exported, package_path=model_file, inductor_configs=lower_configs + ) # Charge-state fold. Present only for a compressed charge-conditioned # descriptor, whose frozen tables the deployment rebuilds once when the # state becomes known. charge_state_bytes: bytes | None = None if charge_state_descriptor is not None: + long_range = _charge_state_long_range(data, metadata) metadata["charge_state_constants"] = _match_charge_state_constants( - charge_state_descriptor, exported + charge_state_descriptor, exported, long_range ) charge_state_bytes = _compile_charge_state_fold( - charge_state_descriptor, aoti_configs + charge_state_descriptor, aoti_configs, long_range ) # Second artifact: with-comm. Only for descriptors whose message diff --git a/deepmd/pt_expt/utils/tabulate_ops.py b/deepmd/pt_expt/utils/tabulate_ops.py index 3e2f3db13b..5fedf66910 100644 --- a/deepmd/pt_expt/utils/tabulate_ops.py +++ b/deepmd/pt_expt/utils/tabulate_ops.py @@ -50,6 +50,12 @@ def _try_register_fake(op_name: str, fn: Callable[..., Any]) -> None: """Register a fake implementation, silently skipping if already registered.""" if op_name in _registered: return + # CompositeImplicitAutograd operators use their decomposition on fake + # tensors; PyTorch forbids a separate Meta registration for such operators. + if torch._C._dispatch_has_kernel_for_dispatch_key( + op_name, "CompositeImplicitAutograd" + ): + return try: torch.library.register_fake(op_name)(fn) _registered.add(op_name) diff --git a/deepmd/pt_expt/utils/vesin_graph_builder.py b/deepmd/pt_expt/utils/vesin_graph_builder.py index cc999bd1af..9e9b73defb 100644 --- a/deepmd/pt_expt/utils/vesin_graph_builder.py +++ b/deepmd/pt_expt/utils/vesin_graph_builder.py @@ -1,18 +1,16 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Carry-all NeighborGraph builder backed by vesin.torch (O(N) cell list). - -World-2 counterpart of vesin_neighbor_list.py: instead of building the dense -quartet, it returns per-frame local (i, j, S), then delegates to the array-API -``neighbor_graph_from_ijs`` (which recomputes ``edge_vec`` differentiably from -the ORIGINAL grad-carrying coords). torch-only => lives in pt_expt. - -Scope note: ``vesin.torch``'s API is single-system, so this builder LOOPS over -frames in Python (~1 ms/frame call overhead measured on GPU). It is intended -for ``nf == 1`` inference and CPU use. Inference ``neighbor_graph_method="auto"`` -(:func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`) selects -vesin only when ``nf == 1`` and ``vesin.torch`` is importable (and ``nv`` is -unavailable on CUDA); multi-frame batches stay on ``nv``/``dense``. Training -auto never selects vesin. Prefer ``nv`` (:mod:`.nv_graph_builder`) for batched +"""Vesin graph builders for differentiable and compact deployment inputs. + +Generic graphs recompute displacements from local neighbor pairs and periodic +image shifts, preserving coordinate and cell derivatives. Compact canonical +graphs use Vesin's detached displacements for exported models with explicit +forces. Both layouts retain every active neighbor within the exact cutoff. + +Vesin's API is single-system, so this builder loops over frames in Python. +It is intended for ``nf == 1`` inference and CPU use. Inference ``neighbor_graph_method="auto"`` +(:func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`) prefers +vesin for single-frame CUDA inputs; multi-frame batches stay on ``nv``/``dense``. +Training auto never selects vesin. Prefer ``nv`` (:mod:`.nv_graph_builder`) for batched multi-frame GPU work, which batches all frames in one kernel. """ @@ -40,57 +38,16 @@ attach_edge_csr, neighbor_graph_from_ijs, ) +from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + canonical_graph_from_pairs, +) from deepmd.pt_expt.utils.vesin_neighbor_list import ( is_vesin_torch_available, ) - - -def vesin_search_ijs( - positions: torch.Tensor, - cell: torch.Tensor | None, - periodic: bool, - rcut: float, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Raw ``vesin.torch`` neighbor search returning ``(ii, jj, ss)`` as int64. - - The caller is responsible for ensuring ``vesin.torch`` is importable (check - :func:`~deepmd.pt_expt.utils.vesin_neighbor_list.is_vesin_torch_available` - before calling). ``positions`` must be detached (the search is - non-differentiable). - - Parameters - ---------- - positions : (nloc, 3) local-frame coordinates, already detached. - cell : (3, 3) box matrix for periodic systems, or ``None`` for non-periodic. - For non-periodic systems a zero box is constructed internally. - periodic : whether the system is periodic. - rcut : neighbor cutoff radius. - device : device to pin as the ambient default. ``vesin.torch`` allocates - some internal tensors on the ambient default device, which may be a - fake/other device in some test contexts (e.g. a placeholder CUDA - default); pinning it here prevents spurious CUDA initializations. - - Returns - ------- - ii : (E,) int64 center local indices. - jj : (E,) int64 neighbor local indices. - ss : (E, 3) int64 periodic image shifts. - """ - import vesin.torch as _vesin_torch - - box = ( - cell if periodic else torch.zeros((3, 3), dtype=positions.dtype, device=device) - ) - nl = _vesin_torch.NeighborList(cutoff=float(rcut), full_list=True) - with torch.device(device): - ii, jj, ss = nl.compute( - points=positions, - box=box, - periodic=periodic, - quantities="ijS", - ) - return ii.to(torch.int64), jj.to(torch.int64), ss.to(torch.int64).reshape(-1, 3) +from deepmd.pt_expt.utils.vesin_search import ( + VesinNeighborSearch, +) def build_neighbor_graph_vesin( @@ -104,6 +61,7 @@ def build_neighbor_graph_vesin( canonicalize: bool = False, pair_excl: PairExcludeMask | None = None, compact: bool = False, + search: VesinNeighborSearch | None = None, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using vesin.torch's O(N) cell list. @@ -135,6 +93,9 @@ def build_neighbor_graph_vesin( compact Passed to :func:`apply_pair_exclusion`; see that function for details. Ignored when ``pair_excl`` is ``None``. + search + Optional persistent search shared across evaluations. Its skin only + controls candidate reuse; the returned graph retains the exact cutoff. Returns ------- @@ -173,45 +134,44 @@ def build_neighbor_graph_vesin( canonicalize=canonicalize, ) + if search is None: + search = VesinNeighborSearch( + sorted=(with_csr or canonicalize) and torch.device(dev).type != "cuda" + ) i_parts, j_parts, S_parts, nf_parts = [], [], [], [] for f in range(nf): - pts = coord[f].detach() - cell_f = box[f].detach() if periodic else None - ii, jj, ss = vesin_search_ijs(pts, cell_f, periodic, rcut, dev) + cell_f = box[f] if periodic else None + ii, jj, ss = search.compute(coord[f], cell_f, rcut, frame=f) i_parts.append(ii) j_parts.append(jj) S_parts.append(ss) nf_parts.append(torch.full((ii.shape[0],), f, dtype=torch.int64, device=dev)) - # guard torch.cat against empty part lists (nf == 0), mirroring ase_builder - i_all = ( - torch.cat(i_parts) - if i_parts - else torch.zeros((0,), dtype=torch.int64, device=dev) - ) - j_all = ( - torch.cat(j_parts) - if j_parts - else torch.zeros((0,), dtype=torch.int64, device=dev) - ) - S_all = ( - torch.cat(S_parts) - if S_parts - else torch.zeros((0, 3), dtype=torch.int64, device=dev) - ) - nf_all = ( - torch.cat(nf_parts) - if nf_parts - else torch.zeros((0,), dtype=torch.int64, device=dev) - ) + # A single frame already owns contiguous edge arrays. Empty batches retain + # the same empty-array contract as the backend-independent ASE builder. + if nf == 1: + i_all, j_all, S_all, nf_all = i_parts[0], j_parts[0], S_parts[0], nf_parts[0] + elif nf > 1: + i_all = torch.cat(i_parts) + j_all = torch.cat(j_parts) + S_all = torch.cat(S_parts) + nf_all = torch.cat(nf_parts) + else: + i_all = torch.zeros((0,), dtype=torch.int64, device=dev) + j_all = torch.zeros((0,), dtype=torch.int64, device=dev) + S_all = torch.zeros((0, 3), dtype=torch.int64, device=dev) + nf_all = torch.zeros((0,), dtype=torch.int64, device=dev) # virtual atoms (atype < 0) are excluded as centers AND neighbors — the # World-2 builder contract shared with the dense reference builder; the # geometric search above cannot know about them. at = torch.as_tensor(atype, device=dev).reshape(nf, nloc) keep = (at[nf_all, i_all] >= 0) & (at[nf_all, j_all] >= 0) - i_all, j_all, nf_all = i_all[keep], j_all[keep], nf_all[keep] - S_all = S_all[keep] + indices = torch.nonzero(keep, as_tuple=False).flatten() + i_all = i_all.index_select(0, indices) + j_all = j_all.index_select(0, indices) + nf_all = nf_all.index_select(0, indices) + S_all = S_all.index_select(0, indices) # i = center (dst), j = neighbor (src); pass ORIGINAL coord/box # (grad-carrying). Unlike the nv builder, vesin's cell list handles @@ -231,5 +191,81 @@ def build_neighbor_graph_vesin( at_flat = torch.as_tensor(atype, device=dev).reshape(-1) graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact) if with_csr or canonicalize: - graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize) + graph = attach_edge_csr( + graph, + nf * nloc, + canonicalize=canonicalize, + destination_sorted=search.sorted and (pair_excl is None or compact), + ) return graph + + +def build_canonical_graph_vesin( + coord: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + search: VesinNeighborSearch, + node_mask: torch.Tensor | None = None, +) -> CanonicalGraph: + """Build the compact deployment graph directly from Vesin displacements. + + Parameters + ---------- + coord : torch.Tensor + Coordinates in angstrom with shape (nf, nloc, 3). + box : torch.Tensor or None + Cells in angstrom with shape (nf, 3, 3), or non-periodic boundaries. + rcut : float + Neighbor cutoff radius in angstrom. + search : VesinNeighborSearch + Persistent search with independent candidate storage for each frame. + node_mask : torch.Tensor or None + Active-node mask with shape (nf, nloc), on the coordinate device. + None includes all nodes. Inactive nodes remain in the node axis but + have no incoming or outgoing edges. + + Returns + ------- + CanonicalGraph + Local graph with exact-cutoff float32 geometry and dual CSR. Geometry + is detached; this builder serves exported models with explicit forces. + """ + nf, nloc, _ = coord.shape + centers, sources, vectors = [], [], [] + for frame in range(nf): + if nloc == 0: + break + center, source, edge_vec = search.compute( + coord[frame], + box[frame] if box is not None else None, + rcut, + frame=frame, + return_vectors=True, + ) + if node_mask is not None: + active = node_mask[frame, center] & node_mask[frame, source] + indices = torch.nonzero(active, as_tuple=False).flatten() + center = center.index_select(0, indices) + source = source.index_select(0, indices) + edge_vec = edge_vec.index_select(0, indices) + if frame: + center = center + frame * nloc + source = source + frame * nloc + centers.append(center) + sources.append(source) + vectors.append(edge_vec) + if not centers: + center = torch.empty(0, dtype=torch.int64, device=coord.device) + source = torch.empty_like(center) + edge_vec = coord.new_empty((0, 3)) + elif nf > 1: + center = torch.cat(centers) + source = torch.cat(sources) + edge_vec = torch.cat(vectors) + return canonical_graph_from_pairs( + center, + source, + edge_vec, + torch.full((nf,), nloc, dtype=torch.int64, device=coord.device), + nf * nloc, + ) diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 7453176d32..a2d67b5754 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -36,6 +36,9 @@ edge_schema_from_ij_shifts, merge_frame_edge_schemas, ) +from deepmd.pt_expt.utils.vesin_search import ( + VesinNeighborSearch, +) def is_vesin_torch_available() -> bool: @@ -56,6 +59,9 @@ class VesinNeighborList(NeighborList): namespace and device of the input coordinates. """ + def __init__(self, skin: float = 0.0) -> None: + self._search = VesinNeighborSearch(skin=skin) + def build( self, coord: Any, @@ -119,6 +125,8 @@ def build( atype_t[ff], rcut, sel, + self._search, + ff, ) for ff in range(nframes) ] @@ -143,6 +151,8 @@ def build( atype_t[ff], rcut, sel, + self._search, + ff, ) for ff in range(nframes) ] @@ -193,6 +203,8 @@ def _build_single( atype: torch.Tensor, rcut: float, sel: list[int], + search: VesinNeighborSearch, + frame: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Single-frame ``(i,j,S)`` -> minimal-extended conversion. @@ -220,17 +232,7 @@ def _build_single( cell if periodic else torch.zeros((3, 3), dtype=positions.dtype, device=device) ) - # Delegate the raw search to the shared helper in vesin_graph_builder - # (function-level import: legacy module depends on graph module lazily to - # avoid a module-level cycle — vesin_graph_builder imports - # is_vesin_torch_available from this module). - from deepmd.pt_expt.utils.vesin_graph_builder import ( - vesin_search_ijs, - ) - - ii, jj, ss = vesin_search_ijs( - positions.detach(), cell if periodic else None, periodic, rcut, device - ) + ii, jj, ss = search.compute(positions, cell, rcut, frame=frame) # Phantom atoms (atype < 0) pad a mixed-nloc batch and have no physical # site; the geometric search above cannot know that. Their pairs are # dropped before the ``sum(sel)`` truncation below, so that a phantom can @@ -309,6 +311,8 @@ def _build_single_edges( atype: torch.Tensor, rcut: float, sel: list[int], + search: VesinNeighborSearch, + frame: int, ) -> EdgeNeighborList: """Single-frame ``vesin`` output converted directly to edge vectors.""" device = positions.device @@ -325,14 +329,7 @@ def _build_single_edges( rcut, ) - periodic = cell is not None - from deepmd.pt_expt.utils.vesin_graph_builder import ( - vesin_search_ijs, - ) - - ii, jj, ss = vesin_search_ijs( - positions.detach(), cell if periodic else None, periodic, rcut, device - ) + ii, jj, ss = search.compute(positions, cell, rcut, frame=frame) # ss is int64 from the helper; edge_schema_from_ij_shifts accepts int shifts. return edge_schema_from_ij_shifts( positions=positions, diff --git a/deepmd/pt_expt/utils/vesin_search.py b/deepmd/pt_expt/utils/vesin_search.py new file mode 100644 index 0000000000..04ff5f31dc --- /dev/null +++ b/deepmd/pt_expt/utils/vesin_search.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Persistent geometric neighbor searches shared by inference input schemas.""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + TYPE_CHECKING, +) + +import torch + +if TYPE_CHECKING: + from vesin_torch import ( + NeighborList, + ) + + +class VesinNeighborSearch: + """Reuse Vesin search storage and optional Verlet candidates. + + Each frame owns a calculator. Vesin refreshes displacements and filters + candidates at the requested cutoff on every call, and rebuilds candidates + when motion exceeds half the skin or the cell or point count changes. + Search results are copied so later calls cannot overwrite indices retained + by an earlier autograd graph. + + Parameters + ---------- + skin : float, default: 0.0 + Extra candidate-search radius in angstrom. Zero disables Verlet caching. + sorted : bool, default: False + Whether pairs are returned in ascending center-index order. + """ + + def __init__(self, skin: float = 0.0, *, sorted: bool = False) -> None: + if not math.isfinite(skin) or skin < 0: + raise ValueError( + f"neighbor_skin must be finite and non-negative, got {skin}" + ) + self.skin = float(skin) + self.sorted = sorted + self._key: tuple[float, torch.device] | None = None + self._frames: dict[int, tuple[NeighborList, torch.Tensor | None]] = {} + + def compute( + self, + positions: torch.Tensor, + cell: torch.Tensor | None, + rcut: float, + *, + frame: int = 0, + return_vectors: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return neighbor indices and image shifts or displacement vectors. + + Parameters + ---------- + positions : torch.Tensor + Coordinates in angstrom with shape (N, 3). The search detaches + these coordinates and runs on their device. + cell : torch.Tensor or None + Cell vectors in angstrom with shape (3, 3), with the same dtype + and device as ``positions``. None selects non-periodic boundaries. + rcut : float + Neighbor cutoff radius in angstrom. + frame : int, optional + Non-negative cache slot, by default 0. Each frame in a batch uses + a distinct slot, reused for its successive configurations. + return_vectors : bool, optional + Return neighbor-minus-center vectors in angstrom instead of + lattice shifts. These vectors are detached from the coordinates + and are intended for inference with explicit force outputs. + + Returns + ------- + ii : torch.Tensor + Local center indices with shape (E,), int64. + jj : torch.Tensor + Local neighbor indices with shape (E,), int64. + geometry : torch.Tensor + Integer lattice shifts with shape (E, 3), int64, or displacement + vectors with the dtype of ``positions`` when ``return_vectors`` + is True. All outputs reside on the device of ``positions`` and + retain their values across subsequent searches. + """ + from vesin_torch import ( + NeighborList, + ) + + device = positions.device + key = (float(rcut), device) + if self._key != key: + self._frames.clear() + self._key = key + calculator, previous_cell = self._frames.get(frame, (None, None)) + # Cell deformation can bring uncached periodic images inside the + # cutoff even when Vesin's single-cell displacement bound permits + # reuse. Comparing CUDA cells synchronizes this invalidation decision. + if calculator is None or ( + self.skin > 0 + and ( + (cell is None) != (previous_cell is None) + or ( + cell is not None + and previous_cell is not None + and not torch.equal(cell, previous_cell) + ) + ) + ): + calculator = NeighborList( + cutoff=float(rcut), + full_list=True, + sorted=self.sorted, + skin=self.skin, + ) + # An independent snapshot detects in-place updates to the input. + previous_cell = ( + cell.detach().clone() if self.skin > 0 and cell is not None else None + ) + self._frames[frame] = calculator, previous_cell + box = ( + cell + if cell is not None + else torch.zeros((3, 3), dtype=positions.dtype, device=device) + ) + with torch.device(device): + try: + ii, jj, geometry = calculator.compute( + points=positions.detach(), + box=box.detach(), + periodic=cell is not None, + quantities="ijD" if return_vectors else "ijS", + ) + except RuntimeError as error: + # An empty CUDA Verlet cache can expose a null pair buffer that + # PyTorch cannot wrap. An uncached search supplies valid empty + # tensors; nonempty results and unrelated errors still fail. + if ( + device.type != "cuda" + or self.skin == 0 + or str(error) + != "The specified pointer resides on host memory and is not " + "registered with any CUDA device." + ): + raise + ii, jj, geometry = NeighborList( + cutoff=float(rcut), full_list=True + ).compute( + points=positions.detach(), + box=box.detach(), + periodic=cell is not None, + quantities="ijD" if return_vectors else "ijS", + ) + if ii.numel() != 0: + raise + if not return_vectors: + geometry = geometry.to(torch.int64) + return ii.to(torch.int64), jj.to(torch.int64), geometry.reshape(-1, 3) diff --git a/deepmd/tf/entrypoints/change_bias.py b/deepmd/tf/entrypoints/change_bias.py index b333fc470e..78e2ebd36c 100644 --- a/deepmd/tf/entrypoints/change_bias.py +++ b/deepmd/tf/entrypoints/change_bias.py @@ -47,6 +47,7 @@ ) from deepmd.utils.data_system import ( DeepmdDataSystem, + get_data, ) __all__ = ["change_bias"] @@ -218,13 +219,22 @@ def _change_bias_checkpoint_file( if stop_batch is None and num_epoch is not None: if num_epoch <= 0: raise ValueError("training.num_epoch must be positive.") - # Apply sys_probs and auto_prob from original training config - # to ensure stop_batch calculation matches the original training training_data_config = training_params.get("training_data", {}) - sys_probs = training_data_config.get("sys_probs", None) - auto_prob = training_data_config.get("auto_prob", "prob_sys_size") - data.set_sys_probs(sys_probs=sys_probs, auto_prob_style=auto_prob) - total_numb_batch = compute_total_numb_batch(data.nbatches, data.sys_probs) + if isinstance(training_data_config, list): + # Group probabilities and batch sizes define the training schedule; + # the independent bias dataset has no corresponding group indices. + schedule_data = get_data( + training_data_config, 0, jdata["model"].get("type_map"), None + ) + else: + data.set_sys_probs( + sys_probs=training_data_config.get("sys_probs"), + auto_prob_style=training_data_config.get("auto_prob", "prob_sys_size"), + ) + schedule_data = data + total_numb_batch = compute_total_numb_batch( + schedule_data.nbatches, schedule_data.sys_probs + ) if total_numb_batch <= 0: raise ValueError("Total number of training batches must be positive.") stop_batch = int(np.ceil(num_epoch * total_numb_batch)) diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 068bef876f..e88c69fcf1 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -670,6 +670,11 @@ def __init__( type_map: list[str] | None = None, **kwargs: Any, ) -> None: + if kwargs.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) super().__init__( descriptor=descriptor, fitting=fitting_net, type_map=type_map, **kwargs ) diff --git a/deepmd/tf2/atomic_model/dp_atomic_model.py b/deepmd/tf2/atomic_model/dp_atomic_model.py index 2d68ba81ef..966c005b10 100644 --- a/deepmd/tf2/atomic_model/dp_atomic_model.py +++ b/deepmd/tf2/atomic_model/dp_atomic_model.py @@ -41,6 +41,8 @@ class tf2_atomic_model(dpmodel_atomic_model): """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls = None + """No long-range module is implemented in this backend.""" @tf.autograph.experimental.do_not_convert def make_atom_mask( diff --git a/deepmd/tf2/entrypoints/train.py b/deepmd/tf2/entrypoints/train.py index a026bcc03f..3350e98c1c 100644 --- a/deepmd/tf2/entrypoints/train.py +++ b/deepmd/tf2/entrypoints/train.py @@ -159,7 +159,7 @@ def factory( type_map = list(task_config.model_params.get("type_map", [])) ipt_type_map = type_map if type_map else None train_data = get_data( - dict(task_config.training_data_params), + task_config.training_data_params, None, ipt_type_map, None, @@ -244,7 +244,7 @@ def update_sel( for task_config in iter_training_task_configs(jdata): type_map = task_config.model_params.get("type_map") train_data = get_data( - dict(task_config.training_data_params), + task_config.training_data_params, 0, type_map, None, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 9f9c98c26c..c7929a0334 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -36,6 +36,9 @@ DEFAULT_FRACTION_INNER, DEFAULT_FRACTION_OUTER, ) +from deepmd.utils.data_config import ( + data_group_configs, +) from deepmd.utils.eval_metrics import ( ENERGY_FULL_VALIDATION_PROFILE, FULL_VALIDATION_PROFILES, @@ -234,8 +237,8 @@ def spin_args() -> list[Argument]: "`native` injects the per-atom spin vector as an equivariant feature " "directly into the descriptor and " "derives the magnetic force as the negative spin gradient of the energy, " - "without virtual atoms. `deepspin` uses the classical DeepSpin virtual-atom " - "representation and is the default. Other models always use the `deepspin` scheme." + "without virtual atoms. Set this to `native` for DPA4/SeZM and DPA4C. " + "Other models use the default `deepspin` virtual-atom representation." ) doc_allow_missing_label = ( "Whether to admit training systems that lack a `spin` data file, filling their " @@ -2928,6 +2931,68 @@ def descrpt_variant_type_args(exclude_hybrid: bool = False) -> Variant: fitting_args_plugin = ArgsPlugin() +@fitting_args_plugin.register("density", doc=supported_backends("pt")) +def fitting_density() -> list[Argument]: + doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." + doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." + doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." + doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' + doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_trainable = "Whether the parameters in the fitting net are trainable. This option can be\n\n\ +- bool: True if all parameters of the fitting net are trainable, False otherwise.\n\n\ +- list of bool: Specifies if each layer is trainable. Since the fitting net is composed by hidden layers followed by a output layer, the length of this list should be equal to len(`neuron`)+1." + doc_seed = "Random seed for parameter initialization of the fitting net" + + # numb_aparam is deliberately not advertised: the fitting consumes + # grid-point descriptor rows, which have no per-atom parameters. + # rcond is deliberately not advertised: the per-type output shift it + # would control is disabled for density models. + return [ + Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_default_fparam, + ), + Argument( + "dim_case_embd", + int, + optional=True, + default=0, + doc=doc_dim_case_embd, + ), + Argument( + "neuron", + list[int], + optional=True, + default=[120, 120, 120], + alias=["n_neuron"], + doc=doc_neuron, + ), + Argument( + "activation_function", + str, + optional=True, + default="tanh", + doc=doc_activation_function, + ), + Argument("precision", str, optional=True, default="default", doc=doc_precision), + Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), + Argument( + "trainable", + [list[bool], bool], + optional=True, + default=True, + doc=doc_trainable, + ), + Argument("seed", [int, None], optional=True, doc=doc_seed), + ] + + @fitting_args_plugin.register( "ener", doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") + doc_ener, @@ -3546,6 +3611,310 @@ def fitting_variant_type_args() -> Variant: ) +# --- Long-range module configurations: --- # +long_range_args_plugin = ArgsPlugin() +interaction_kernel_args_plugin = ArgsPlugin() + + +@interaction_kernel_args_plugin.register( + "les", + doc="The latent Ewald summation kernel erf(alpha r) / r, the interaction of " + "two Gaussian charge clouds: it reaches 1/r beyond a few 1/alpha and stays " + "finite at contact.", +) +def interaction_kernel_les() -> list[Argument]: + doc_alpha = "The inverse width alpha of the kernel, in Å^-1. Must be positive." + doc_trainable = "Whether alpha is trained." + return [ + Argument("alpha", float, optional=True, default=1.0, doc=doc_alpha), + Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), + ] + + +@interaction_kernel_args_plugin.register( + "sog", + doc="The sum-of-Gaussians kernel sum_k a_k exp(-r^2 / (2 sigma_k^2)).", +) +def interaction_kernel_sog() -> list[Argument]: + doc_amplitude = ( + "The amplitudes a_k of the Gaussians, in Å^-1, of either sign: one value " + "shared by every Gaussian or one value per Gaussian. Required with " + "`width`; with the default widths it defaults to " + "sqrt(2 / pi) ln(`width_ratio`) / sigma_k, the weights of a discretized " + "1 / r, so that the default kernel follows 1 / r, less a nearly " + "constant offset, from a few `base_width` to a fraction of the widest " + "Gaussian. The offset exerts no force between fixed charges; through " + "charges that respond to the positions it does." + ) + doc_width = ( + "The widths sigma_k of the Gaussians, in Å, each positive. Defaults to " + "the geometric sequence `base_width` * `width_ratio`^k for " + "k = 0, ..., `n_gaussian` - 1." + ) + doc_n_gaussian = "The number of Gaussians of the default widths." + doc_base_width = "The width of the first default Gaussian, in Å." + doc_width_ratio = "The ratio of consecutive default widths. Must exceed one." + doc_trainable = "Whether the amplitudes and widths are trained." + return [ + Argument( + "amplitude", + [float, list[float], type(None)], + optional=True, + default=None, + doc=doc_amplitude, + ), + Argument( + "width", + [list[float], type(None)], + optional=True, + default=None, + doc=doc_width, + ), + Argument("n_gaussian", int, optional=True, default=12, doc=doc_n_gaussian), + Argument( + "base_width", + float, + optional=True, + default=2.180230445405648, + doc=doc_base_width, + ), + Argument( + "width_ratio", + float, + optional=True, + default=1.6297670882677647, + doc=doc_width_ratio, + ), + Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), + ] + + +def latent_charge_net_args() -> list[Argument]: + doc_neuron = "The number of neurons in each hidden layer of the charge network." + doc_activation_function = f"The activation function in the charge network. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Defaults to the activation function of the energy fitting." + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection. Defaults to the setting of the energy fitting.' + doc_precision = f"The precision of the charge network and kernel parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Defaults to the precision of the energy fitting." + doc_trainable = "Whether the parameters of the charge network are trained. A list sets every hidden layer and the output layer separately." + doc_seed = "Random seed for parameter initialization of the charge network. Defaults to a seed derived from the seed of the energy fitting." + return [ + Argument("neuron", list[int], optional=True, default=[64, 64], doc=doc_neuron), + Argument( + "activation_function", str, optional=True, doc=doc_activation_function + ), + Argument("resnet_dt", bool, optional=True, doc=doc_resnet_dt), + Argument("precision", str, optional=True, doc=doc_precision), + Argument( + "trainable", + [list[bool], bool], + optional=True, + default=True, + doc=doc_trainable, + ), + Argument("seed", [int, None], optional=True, doc=doc_seed), + ] + + +def charge_constraint_args() -> list[Argument]: + doc_source = ( + "The source of the total charge of a frame: `neutral` for zero, " + "`charge_spin` for the charge of the frame charge/spin condition, which " + "requires a descriptor with `add_chg_spin_ebd`, or `fparam` for the frame " + "parameter at `fparam_index`." + ) + doc_fparam_index = "The column of the frame parameters holding the total charge, used when `source` is `fparam`." + doc_channel = "The latent charge channel that is constrained to the total charge; the other channels stay unconstrained." + return [ + Argument( + "source", + str, + optional=True, + default="neutral", + doc=doc_source, + extra_check=lambda x: x in ("neutral", "charge_spin", "fparam"), + extra_check_errmsg="must be one of 'neutral', 'charge_spin' or 'fparam'", + ), + Argument("fparam_index", int, optional=True, default=0, doc=doc_fparam_index), + Argument("channel", int, optional=True, default=0, doc=doc_channel), + ] + + +@long_range_args_plugin.register( + "latent_charge", + doc="Interaction of latent charges predicted from the descriptor. " + "For a non-periodic system the kernel is summed directly over every pair " + "of atoms; for a periodic system (``periodic: true``) it is summed over " + "the images of the cell in reciprocal space. Every channel interacts " + "through the same kernel; every atom carries half of the interaction of " + "its charges with all others.", +) +def long_range_latent_charge() -> list[Argument]: + doc_channels = "The number of latent charge channels predicted per atom." + doc_charge_net = ( + "The network predicting the latent charges, in units of the elementary " + "charge, from the descriptor; it takes the frame and atomic parameters " + "of the energy fitting." + ) + doc_kernel = "The pair interaction kernel of the latent charges, in Å^-1." + doc_charge_constraint = ( + "Shift one latent charge channel uniformly over the atoms of every frame " + "so that it sums to the total charge of the frame. Unconstrained if not " + "set." + ) + doc_periodic = ( + "Evaluate the module on periodic systems as well: the kernel is summed " + "over the images of the cell in reciprocal space, with the " + "zero-wave-number mode dropped (the tin-foil convention, which places " + "a uniform neutralizing background under a channel with a net " + "charge), less the interaction of every charge with itself. " + "Non-periodic systems keep the direct sum." + ) + doc_kspace_tolerance = ( + "The relative weight of the slowest-decaying kernel mode at the " + "wave-number cutoff k_max of a periodic system." + ) + doc_k_grid_mode = ( + "How training evaluates the reciprocal sum of a periodic batch: " + "'pme', the default of a periodic module, by the smooth particle-mesh " + "Ewald method on a mesh sized by the cells of the batch at " + "`pme_spacing`, whose cost grows nearly linearly with the system; " + "'ragged' over the wave numbers " + "of every frame within k_max, enumerated per batch, so that the k axis " + "is a dynamic dimension; or 'fixed' over one integer grid shared by " + "the frames, sized by the largest cell of the data statistics or by " + "`k_grid_bounds`, with the points a frame does not reach masked, so " + "that the k axis has a static length (a frame whose cell needs points " + "beyond the grid is truncated to it, with a warning). A frozen model " + "evaluates every periodic frame on the mesh of its cell, whatever the " + "mode. 'fixed' and 'pme' require `periodic: true`." + ) + doc_k_grid_bounds = ( + "The per-axis integer bounds [n1, n2, n3] of the fixed k grid, in " + "place of the data statistics; only meaningful with " + "`k_grid_mode: fixed`." + ) + doc_decouple_k_grid = ( + "Decouple the wave-number cutoff of a periodic system from the kernel " + "parameters: it is derived once from the initial parameters and kept " + "for the life of the model, so that the k grid and the mesh stay put " + "while the kernel trains (the kernel values still follow the live " + "parameters). The frozen cutoff is serialized with the model. " + "Requires `periodic: true`." + ) + doc_pme_spacing = ( + "The largest spacing of the mesh of a periodic frame, in Angstrom; " + "the default is three quarters of the Nyquist spacing of the " + "wave-number cutoff, 0.75 pi / k_max. Requires `periodic: true`." + ) + doc_pme_order = ( + "The order of the B-splines that spread the charges to the mesh, 8 " + "by default. Requires `periodic: true`." + ) + return [ + Argument("channels", int, optional=True, default=1, doc=doc_channels), + Argument( + "charge_net", + dict, + latent_charge_net_args(), + optional=True, + doc=doc_charge_net, + ), + Argument( + "kernel", + dict, + [], + [ + Variant( + "type", + interaction_kernel_args_plugin.get_all_argument(), + optional=True, + default_tag="les", + doc="The type of the kernel.", + ) + ], + optional=True, + default={"type": "les"}, + doc=doc_kernel, + ), + Argument( + "charge_constraint", + dict, + charge_constraint_args(), + optional=True, + default=None, + doc=doc_charge_constraint, + ), + Argument("periodic", bool, optional=True, default=False, doc=doc_periodic), + Argument( + "kspace_tolerance", + float, + optional=True, + default=1e-8, + doc=doc_kspace_tolerance, + ), + Argument( + "k_grid_mode", + [str, None], + optional=True, + default=None, + doc=doc_k_grid_mode, + extra_check=lambda x: x in (None, "ragged", "fixed", "pme"), + extra_check_errmsg="must be 'ragged', 'fixed' or 'pme'", + ), + Argument( + "k_grid_bounds", + list[int], + optional=True, + default=None, + doc=doc_k_grid_bounds, + ), + Argument( + "decouple_k_grid", + bool, + optional=True, + default=False, + doc=doc_decouple_k_grid, + ), + Argument( + "pme_spacing", + [float, None], + optional=True, + default=None, + doc=doc_pme_spacing, + ), + Argument( + "pme_order", + [int, None], + optional=True, + default=None, + doc=doc_pme_order, + ), + ] + + +def long_range_args() -> Argument: + doc_long_range = ( + "A long-range module added to the atomic energy of the energy fitting. " + "It reads the same descriptor and is evaluated on complete systems: " + "non-periodic ones directly, periodic ones in reciprocal space when " + "the module serves them. Currently validated with the DPA4C " + "descriptor." + ) + return Argument( + "long_range", + dict, + [], + [ + Variant( + "type", + long_range_args_plugin.get_all_argument(), + doc="The type of the long-range module.", + ) + ], + optional=True, + doc=supported_backends("pt_expt") + doc_long_range, + ) + + # --- Modifier configurations: --- # modifier_args_plugin = ArgsPlugin() @@ -3643,7 +4012,7 @@ def model_args( "For an energy model it is the energy of an isolated atom, so with `vacuum_ref` in the fitting net an atom " "without neighbors gives exactly this energy. Four forms are accepted: a dict keyed by element symbol, " "e.g. `{'energy': {'O': -430.1, 'H': -13.6}}`; the name of a bundled table of isolated-atom energies, e.g. " - "`{'energy': 'omat24'}`, one of `mptraj`, `omat24`, `omol25`, `omc25` and `odac25`, which takes precedence " + "`{'energy': 'omat24'}`, one of `mptraj`, `omat24`, `omol25`, `omc25`, `odac25`, `openmd` and `domain_cluster`, which takes precedence " "over a file of the same name; " "the path of a JSON file holding such a dict, e.g. `{'energy': 'e0.json'}`, relative to the working directory; " "or a list with one entry per type of the `type_map` (`null` leaves a type unassigned), which is also the form " @@ -3651,7 +4020,9 @@ def model_args( "Elements outside the `type_map` are ignored and every element that occurs in the data must be assigned. " "A table is resolved once when the input is processed and its values are stored in the model. An assigned " "output is taken from the preset without statistics, both when a model is initialized and when its bias is " - "changed by fine-tuning or `dp change-bias`. Dipole models apply no output bias and take no preset." + "changed by fine-tuning. `dp change-bias` fits from its input data without reapplying the model's stored " + "preset; explicit `--bias-value` values override the bias directly. " + "Dipole models apply no output bias and take no preset." ) doc_finetune_head = ( "The chosen fitting net to fine-tune on, when doing multi-task fine-tuning. " @@ -3843,6 +4214,7 @@ def standard_model_args() -> Argument: default={}, doc=supported_backends("pt", "jax", "pd", "pt_expt", "tf2") + doc_info, ), + long_range_args(), *_bridging_method_args(), ], doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") @@ -5183,6 +5555,28 @@ def loss_ener() -> list[Argument]: ] +@loss_args_plugin.register("grid_density", doc=supported_backends("pt")) +def loss_grid_density() -> list[Argument]: + doc_start_pref_d = start_pref("density", abbr="d") + doc_limit_pref_d = limit_pref("density") + return [ + Argument( + "start_pref_d", + [float, int], + optional=True, + default=1.00, + doc=doc_start_pref_d, + ), + Argument( + "limit_pref_d", + [float, int], + optional=True, + default=1.00, + doc=doc_limit_pref_d, + ), + ] + + @loss_args_plugin.register("dens", doc=supported_backends("pt")) def loss_dens() -> list[Argument]: doc_start_pref_e = start_pref("energy", abbr="e") @@ -5773,15 +6167,15 @@ def loss_args() -> list[Argument]: # --- Training configurations: --- # -def training_data_args() -> list[ - Argument -]: # ! added by Ziyao: new specification style for data systems. +def training_data_args() -> Argument: link_sys = make_link("systems", "training/training_data/systems") doc_systems = ( "The data systems for training. " "This key can be a list or a str. " "When provided as a string, it can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. " - "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories." + "When provided as a list, each string item is expanded recursively in input order. " + "The PyTorch and PyTorch Exportable backends also recognize LMDB datasets, " + "either directly or beneath a parent directory." ) doc_patterns = ( "The customized patterns used in `rglob` to collect all training systems. " @@ -5797,10 +6191,15 @@ def training_data_args() -> list[ - string "mix:N": LMDB data only. Frames of different atom counts share a batch, filled until the next frame would push the atom axis of the batch past `N`. How that axis is laid out follows from the model and needs no configuration of its own: a descriptor reading a flat node axis (the graph route of the PyTorch Exportable backend) takes the frames of a batch concatenated, so `N` counts real atoms and nothing is padded; every other descriptor takes them padded to the widest frame of the batch, so `N` counts the padded slots `nframes * max_nloc` and the shorter frames carry phantom atoms with `atype = -1` that the neighbor list, the model and the loss all skip. Unlike `"max:N"`, which leaves an under-filled batch whenever an nloc group is small, this keeps every batch close to `N` atoms. A lone frame with nloc > `N` still forms a batch of its own.\n\n\ If MPI is used, the value should be considered as the batch size per task.' doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\ -- "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ +- "prob_uniform" : each retained system supplies the same long-run share of sampled frames. This includes internal LMDB systems and corrects for differing batch frame counts.\n\n\ - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\ - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is divided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.' - doc_sys_probs = "A list of float if specified. Should be of the same length as `systems`, specifying the probability of each system." + doc_sys_probs = ( + "A list of float if specified. Should be of the same length as `systems`, " + "specifying the probability of each system. A group expanding to multiple " + "sources including LMDB must omit this option; use separate training_data " + "groups with prob to assign source probabilities." + ) doc_min_pair_dist = ( "Minimum pairwise atomic distance threshold in Å. " "Frames containing any atom pair closer than this distance are excluded " @@ -5842,7 +6241,10 @@ def training_data_args() -> list[ str, optional=True, default="prob_sys_size", - doc=doc_auto_prob_style, + doc=doc_auto_prob_style + + " When a group expands to multiple sources including LMDB, only " + "prob_sys_size and prob_uniform are supported. Use separate groups to configure " + "individual LMDB system probabilities.", alias=[ "auto_prob_style", ], @@ -5862,12 +6264,28 @@ def training_data_args() -> list[ default=0.0, doc=supported_backends("pt", "pt_expt") + doc_min_pair_dist, ), + Argument( + "prob", + [float, None], + optional=True, + default=None, + doc="The probability of selecting a training batch from this entire " + "group. Applies when training_data is a list of group dictionaries. " + "Assigned probabilities must be non-negative and sum to at most one. " + "Groups without prob share the remaining probability in proportion " + "to their batch counts. If every group specifies prob, their sum " + "must be one. Within-group auto_prob and sys_probs remain independent.", + ), ] - doc_training_data = "Configurations of training data." + doc_training_data = ( + "Training data configured as one dictionary or a nonempty list of group " + "dictionaries. Each group has its own systems, batch_size and within-group " + "sampling rules; prob optionally controls the group's total probability." + ) return Argument( "training_data", - dict, + [dict, list[dict]], optional=True, sub_fields=args, sub_variants=[], @@ -5884,7 +6302,9 @@ def validation_data_args() -> list[ "The data systems for validation. " "This key can be a list or a str. " "When provided as a string, it can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. " - "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories." + "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories. " + "The PyTorch and PyTorch Exportable backends also recognize LMDB datasets, " + "including multiple LMDB paths and parent directories containing LMDB datasets." ) doc_patterns = ( "The customized patterns used in `rglob` to collect all validation systems. " @@ -5898,10 +6318,14 @@ def validation_data_args() -> list[ - string "filter:N": the same as `"max:N"` but additionally drops data whose atom count exceeds `N`. For npy data this removes whole systems with natoms > `N`; for LMDB data this removes individual frames with nloc > `N`.\n\n\ - string "mix:N": LMDB data only. Frames of different atom counts share a batch, filled until the next frame would push the atom axis of the batch past `N`. How that axis is laid out follows from the model: a descriptor reading a flat node axis takes the frames of a batch concatenated, so `N` counts real atoms and nothing is padded; every other descriptor takes them padded to the widest frame of the batch, so `N` counts the padded slots `nframes * max_nloc` and the shorter frames carry phantom atoms that the neighbor list, the model and the loss all skip.' doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\ -- "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ +- "prob_uniform" : each retained system supplies the same long-run share of sampled frames. This includes internal LMDB systems and corrects for differing batch frame counts.\n\n\ - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\ - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is divided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.' - doc_sys_probs = "A list of float if specified. Should be of the same length as `systems`, specifying the probability of each system." + doc_sys_probs = ( + "A list of float if specified. Should be of the same length as `systems`, " + "specifying the probability of each system. Must be omitted when the " + "validation data expands to multiple sources including LMDB." + ) doc_numb_btch = "An integer that specifies the number of batches to be sampled for each validation period." args = [ @@ -5928,7 +6352,9 @@ def validation_data_args() -> list[ str, optional=True, default="prob_sys_size", - doc=doc_auto_prob_style, + doc=doc_auto_prob_style + + " Only prob_sys_size and prob_uniform are supported when systems expands to " + "multiple sources including LMDB.", alias=[ "auto_prob_style", ], @@ -6039,6 +6465,16 @@ def training_args( doc_disp_freq = "The frequency of printing learning curve." doc_save_freq = "The frequency of saving check point." doc_save_ckpt = "The path prefix of saving check point files." + doc_reset_data_progress = ( + "On --restart, discard the saved training-data cursors and task/system " + "selection RNG, retaining their newly initialized state. Model weights, " + "optimizer, learning-rate schedule and training step still resume. " + "Use this explicit reset when relocating data or intentionally changing " + "data settings; batches may repeat and exact data-order continuation " + "is unavailable. By default, saved data progress is restored strictly, " + "including dataset locations. Has no effect on fresh training, " + "initialization, finetuning or checkpoints without data progress." + ) doc_save_dir = ( "The directory in which periodic checkpoint files are written, " "including the regular checkpoints (the `save_ckpt` prefix) and, when " @@ -6126,8 +6562,11 @@ def training_args( doc_model_prob = ( "The visiting probability of each model for each training step in the " "multi-task mode. Only used when num_epoch_dict is not set. If not set " - "or an empty dict, defaults to weights proportional to the number of " - "systems per task." + "or an empty dict, defaults to each task's native data size: the number " + "of NPY/HDF5 systems, or retained LMDB frames on the PyTorch and " + "PyTorch Exportable backends. Grouped inputs sum their native sizes " + "without applying group probabilities. Batch size and nominal epoch " + "length do not affect these default weights." ) doc_data_dict = "The multiple definition of the data, used in the multi-task mode." doc_acc_freq = "Gradient accumulation steps (number of steps to accumulate gradients before performing an update)." @@ -6237,6 +6676,13 @@ def training_args( Argument( "save_ckpt", str, optional=True, default="model.ckpt", doc=doc_save_ckpt ), + Argument( + "reset_data_progress", + bool, + optional=True, + default=False, + doc=supported_backends("pt_expt") + doc_reset_data_progress, + ), Argument("max_ckpt_keep", int, optional=True, default=5, doc=doc_max_ckpt_keep), Argument( "ckpt_keep_ratio", @@ -6883,6 +7329,7 @@ def gen_json_schema(multi_task: bool = False) -> str: JSON schema. """ from deepmd.utils.json_schema import ( + with_data_groups, with_model_presets, ) @@ -6894,7 +7341,9 @@ def gen_json_schema(multi_task: bool = False) -> str: ) return json.dumps( with_model_presets( - generate_json_schema(arg), generate_json_schema(model_args()), multi_task + with_data_groups(generate_json_schema(arg), multi_task), + generate_json_schema(model_args()), + multi_task, ) ) @@ -6967,17 +7416,86 @@ def validate_no_multitask_lora(data: dict[str, Any], multi_task: bool = False) - ) +def _apply_density_env_protection_default(data: dict[str, Any]) -> None: + """Default env_protection to 1e-6 for density models. + + Applied at normalization time so that the recorded model_def_script and + the built model agree on this field. Grid points may legitimately + coincide with atoms, and the default 0.0 would let the 1/r terms in the + environment matrix produce NaN densities. 1e-6 is chosen as the minimal + perturbation that keeps those terms finite (order 1e6 at exact + coincidence) without shifting normal environments measurably; users can + still set any positive value (e.g. 0.1) explicitly. An explicit + ``env_protection: 0.0`` cannot be distinguished from the schema default + (which normalizes to 0.0) and receives the same override with a warning. + """ + + def _fix(model: dict[str, Any]) -> None: + if model.get("fitting_net", {}).get("type") != "density": + return + descriptor = model.get("descriptor", {}) + fill = 1e-6 + # a hybrid descriptor has no top-level env_protection; each + # sub-descriptor carries its own, and the model build requires all + # of them to agree. Filling each omitted sub-descriptor with 1e-6 + # while a sibling carries an explicit value (e.g. 0.1) would produce + # a mixed list that the build then rejects, so omitted entries take + # the siblings' common explicit value instead. + if descriptor.get("type") == "hybrid": + sub_descriptors = descriptor.get("list", []) + explicit = {sub.get("env_protection", 0.0) for sub in sub_descriptors} + explicit.discard(0.0) + if len(explicit) == 1: + fill = explicit.pop() + else: + sub_descriptors = [descriptor] + for sub in sub_descriptors: + if sub.get("env_protection", 0.0) == 0.0: + log.warning( + "env_protection is 0.0 for a density model; grid points " + "coincident with atoms would produce NaN densities. " + "Setting env_protection to %s.", + fill, + ) + sub["env_protection"] = fill + + model = data.get("model", {}) + if "model_dict" in model: + for sub_model in model["model_dict"].values(): + _fix(sub_model) + else: + _fix(model) + + def normalize( data: dict[str, Any], multi_task: bool = False, *, check: bool = True ) -> dict[str, Any]: """Normalize config values and optionally run strict schema checks.""" base = Argument("base", dict, gen_args(multi_task=multi_task)) data = base.normalize_value(data, trim_pattern="_*") + training = data.get("training", {}) + task_configs = training.get("data_dict", {}).values() if multi_task else (training,) + group_argument = training_data_args() + for task in task_configs: + groups = task.get("training_data") + if groups is None: + continue + data_group_configs(groups) + if isinstance(groups, list): + normalized = [ + group_argument.normalize_value(group, trim_pattern="_*") + for group in groups + ] + if check: + for group in normalized: + group_argument.check_value(group, strict=True) + task["training_data"] = normalized if check: base.check_value(data, strict=True) validate_full_validation_config(data, multi_task=multi_task) _check_dpa3_chg_spin_migration(data) validate_no_multitask_lora(data, multi_task=multi_task) + _apply_density_env_protection_default(data) return data diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index f92dca7c77..f8c62f48e4 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -10,7 +10,6 @@ ) from typing import ( Any, - Protocol, ) import array_api_compat @@ -26,19 +25,15 @@ class RetrySignal(Exception): """Signal to retry execution after OOM error.""" -class BatchSizeBudget(Protocol): - """Limit atom budgets and observe completed batch attempts.""" - - def limit(self, batch_size: int) -> int: - """Return an admissible atom budget no larger than the proposal.""" - - def observe(self, nframes: int) -> None: - """Record an attempt; zero frames denotes an unsuccessful call.""" - - class AutoBatchSize(ABC): - """This class allows DeePMD-kit to automatically decide the maximum - batch size that will not cause an OOM error. + """Grow inference batches and reuse the learned out-of-memory bound. + + Successful full batches grow geometrically. After OOM, growth beyond the + failed bound resumes after two successful full batches; repeated recovery probes + extend that interval to four, then eight batches. This deliberately favors + recovering capacity over permanently avoiding OOM. Short tails and changes + in the number of atoms do not reset the learned state. No allocator profiling + or exact maximum search is needed. Notes ----- @@ -79,15 +74,16 @@ def __init__( self.silent = silent self.current_batch_size = initial_batch_size DP_INFER_BATCH_SIZE = int(os.environ.get("DP_INFER_BATCH_SIZE", 0)) - self._fixed_batch_size = DP_INFER_BATCH_SIZE > 0 - self._budget: BatchSizeBudget | None = None + self._auto_growth = DP_INFER_BATCH_SIZE <= 0 and self.is_gpu_available() + self._successful_batches = 0 + self._reprobe_interval = 2 if DP_INFER_BATCH_SIZE > 0: self.current_batch_size = DP_INFER_BATCH_SIZE self.maximum_working_batch_size = DP_INFER_BATCH_SIZE self.minimal_not_working_batch_size = self.maximum_working_batch_size + 1 else: self.maximum_working_batch_size = initial_batch_size - if self.is_gpu_available(): + if self._auto_growth: self.minimal_not_working_batch_size = 2**31 else: self.minimal_not_working_batch_size = ( @@ -111,7 +107,7 @@ def execute( *, max_nframes: int | None = None, ) -> tuple[int, Any]: - """Execute one batch within the search and backend resource budgets. + """Execute one batch, growing on success and backing off on OOM. Parameters ---------- @@ -140,24 +136,16 @@ def execute( if max_nframes is not None and max_nframes <= 0: raise ValueError("max_nframes must be positive") units_per_frame = max(natoms, 1) - proposed_size = self.current_batch_size - budget = self._get_batch_budget(natoms, max_nframes) - if budget is not None: - if budget is not self._budget: - self._budget = budget - self.maximum_working_batch_size = 0 - self.minimal_not_working_batch_size = 2**31 - self._set_batch_size(budget.limit(self.current_batch_size)) batch_nframes = max(self.current_batch_size // units_per_frame, 1) if max_nframes is not None: batch_nframes = min(batch_nframes, max_nframes) - n_batch = 0 try: n_batch, result = callable(batch_nframes, start_index) except Exception as e: if not self.is_oom_error(e): raise attempted_size = batch_nframes * units_per_frame + self._successful_batches = 0 self.minimal_not_working_batch_size = min( self.minimal_not_working_batch_size, attempted_size ) @@ -175,32 +163,31 @@ def execute( if self.oom_retry_mode: raise RetrySignal from e return 0, None - finally: - if budget is not None: - budget.observe(n_batch) n_tot = n_batch * units_per_frame self.maximum_working_batch_size = max(self.maximum_working_batch_size, n_tot) - if ( - n_tot + units_per_frame > self.current_batch_size - and self.current_batch_size * self.factor - < self.minimal_not_working_batch_size - ): - # Calibration limits this attempt without discarding the requested - # atom budget for subsequent batches. - proposed_size = max( - proposed_size, int(self.current_batch_size * self.factor) - ) - self._set_batch_size( - proposed_size if budget is None else budget.limit(proposed_size) - ) + # A successful *actual* batch can disprove an old OOM bound. A + # truncated probe cannot: the proposed size was not actually tested. + if self._auto_growth and n_tot >= self.minimal_not_working_batch_size: + self.minimal_not_working_batch_size = 2**31 + if self._auto_growth and n_tot + units_per_frame > self.current_batch_size: + self._successful_batches += 1 + self._grow_batch_size(n_tot, allow_reprobe=not self.oom_retry_mode) return n_batch, result - def _get_batch_budget( - self, natoms: int, max_nframes: int | None - ) -> BatchSizeBudget | None: - """Return a backend resource budget, or None for OOM-driven sizing.""" - return None + def _grow_batch_size(self, successful_size: int, *, allow_reprobe: bool) -> None: + """Grow below the failed bound, or schedule an allowed recovery probe.""" + # Large mandatory single frames also justify immediate growth. + next_batch_size = int( + max(self.current_batch_size, successful_size) * self.factor + ) + if next_batch_size >= self.minimal_not_working_batch_size: + if not allow_reprobe or self._successful_batches < self._reprobe_interval: + return + # Dampen repeated recovery probes, including oscillating workloads. + self._reprobe_interval = min(2 * self._reprobe_interval, 8) + self._set_batch_size(next_batch_size) + self._successful_batches = 0 def _set_batch_size(self, batch_size: int) -> None: old_batch_size = self.current_batch_size @@ -322,6 +309,15 @@ def concate_result(r: list[Any]) -> Any: r = r[0] else: r = {kk: concate_result(vv) for kk, vv in results.items()} + # Retry-mode callers discard all hook outputs on OOM and restart this + # entire input. Probing a known failure mid-pass could restart forever. + # Defer that probe until the pass is complete, for the next invocation. + if ( + self._auto_growth + and self.oom_retry_mode + and self._successful_batches >= self._reprobe_interval + ): + self._grow_batch_size(0, allow_reprobe=True) return r @abstractmethod @@ -358,6 +354,10 @@ def set_oom_retry_mode(self, enable: bool) -> None: Callers that want all data to be re-executed must catch :class:`RetrySignal` and restart the full evaluation themselves. + Recovery probes beyond a known failed size are deferred until + :meth:`execute_all` has assembled its complete result, for the next + invocation. Direct :meth:`execute` calls have no completion boundary + and therefore do not schedule those probes in retry mode. Parameters ---------- diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index b5bdb2ea94..2fa225cc64 100644 --- a/deepmd/utils/bridging.py +++ b/deepmd/utils/bridging.py @@ -142,6 +142,9 @@ def _is_dpa4_family(sub: dict) -> bool: _LEARNED_CHILD_KEYS = ( "descriptor", "fitting_net", + # A long-range module belongs to the learned energy; the composition then + # rejects it beside the bridging term. + "long_range", "model_branch_alias", "info", "use_compile", diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 4d73e58d87..715640653f 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -4,6 +4,7 @@ import bisect import copy import functools +import hashlib import logging from collections.abc import ( Iterable, @@ -148,6 +149,8 @@ def __init__( # set counters self.set_count = 0 self.iterator = 0 + self._batch_set_size = 0 + self._batch_shuffle_state = None self.shuffle_test = shuffle_test # set modifier self.modifier = modifier @@ -208,6 +211,9 @@ def add( special_shape : str, optional Name of a loader-defined non-standard shape contract. ``"hessian"`` stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame. + ``"frame_major"`` stores one variable-length ``(npoints, ndof)`` + array per frame (e.g. grid/density), kept as a frame-major + ``(nframes, npoints, ndof)`` tensor without any natoms reshaping. length_scale : str, optional Per-element length scale a derived geometric item measures its default against: ``"absolute"`` gives every element the unit radius, @@ -236,6 +242,7 @@ def add( # only non-standard tensors need this extra shape contract. data_config["special_shape"] = special_shape self.data_dict[key] = data_config + self._invalidate_data_cache() return self def reduce(self, key_out: str, key_in: str) -> "DeepmdData": @@ -264,8 +271,16 @@ def reduce(self, key_out: str, key_in: str) -> "DeepmdData": "repeat": 1, "reduce": key_in, } + self._invalidate_data_cache() return self + def _invalidate_data_cache(self) -> None: + """Discard arrays whose data requirements are no longer current.""" + for key in ("batch_set", "_batch_indices", "test_set"): + self.__dict__.pop(key, None) + if hasattr(self, "_modified_frame_cache"): + self._modified_frame_cache.clear() + def get_data_dict(self) -> dict: """Get the `data_dict`.""" return self.data_dict @@ -322,6 +337,30 @@ def get_item_paddle( """ return self.get_single_frame(index, num_worker) + @functools.cached_property + def _sampling_set_nframes(self) -> np.ndarray: + """Count frames in each set after its on-disk copy weights are applied.""" + sizes = [] + for set_dir, nframes in zip( + self.dirs, np.diff([0, *self.prefix_sum]), strict=True + ): + _, copies = self._load_data( + set_dir, "numb_copy", int(nframes), 1, must=False, default=1, dtype=int + ) + sizes.append(int(copies.sum())) + return np.asarray(sizes, dtype=np.int64) + + def get_mean_batch_nframes(self, batch_size: int) -> float: + """Return the initial pass's mean batch size, including undersized sets. + + A set smaller than one batch is returned whole. Larger sets drop their + short tail, as in :meth:`get_batch`. Copy weights are applied before + counting batches. + """ + sizes = self._sampling_set_nframes + counts = np.maximum(1, sizes // batch_size) + return float(np.dot(counts, np.minimum(sizes, batch_size)) / counts.sum()) + def get_batch(self, batch_size: int) -> dict: """Get a batch of data with `batch_size` frames. The frames are randomly picked from the data system. @@ -330,18 +369,26 @@ def get_batch(self, batch_size: int) -> dict: batch_size size of the batch """ - if hasattr(self, "batch_set"): - set_size = self.batch_set["coord"].shape[0] - else: - set_size = 0 - if self.iterator + batch_size > set_size: + if self.iterator + batch_size > self._batch_set_size: self._load_batch_set(self.dirs[self.set_count % self.get_numb_set()]) self.set_count += 1 - set_size = self.batch_set["coord"].shape[0] + elif not hasattr(self, "_batch_indices"): + # Restore the active order only when it is actually needed. An + # exhausted set can advance directly to the next one above. + iterator = self.iterator + set_size = self._batch_set_size + self._load_batch_set( + self.dirs[(self.set_count - 1) % self.get_numb_set()], + shuffle_state=self._batch_shuffle_state, + ) + if self._batch_set_size != set_size: + raise ValueError("The restored data set has a different numb_copy size") + self.iterator = iterator + set_size = self._batch_set_size iterator_1 = self.iterator + batch_size if iterator_1 >= set_size: iterator_1 = set_size - idx = np.arange(self.iterator, iterator_1, dtype=np.int64) + idx = self._batch_indices[self.iterator : iterator_1] self.iterator += batch_size ret = self._get_subdata(self.batch_set, idx) if "pair_margin" in self.data_dict: @@ -353,6 +400,84 @@ def get_batch(self, batch_size: int) -> dict: ) return ret + def _progress_identity(self) -> dict[str, Any]: + """Describe the source using metadata already read at construction.""" + return { + "sets": [str(path) for path in self.dirs], + "prefix_sum": self.prefix_sum, + "atom_types": hashlib.sha256(self.atom_type.tobytes()).hexdigest(), + "type_map": self.type_map, + "sort_atoms": self.sort_atoms, + "mixed_type": self.mixed_type, + "pbc": self.pbc, + } + + def state_dict(self) -> dict[str, Any]: + """Return compact batch progress, without frame arrays or permutations. + + The shared :mod:`deepmd.utils.random` state must be checkpointed + separately, once for all readers. Only the state before the active + set's shuffle is retained here to reconstruct its order lazily. + """ + return copy.deepcopy( + { + "version": 1, + "identity": self._progress_identity(), + "set_count": int(self.set_count), + "iterator": int(self.iterator), + "set_size": int(self._batch_set_size), + "shuffle_state": self._batch_shuffle_state, + } + ) + + def _validate_state_dict(self, state: dict[str, Any]) -> None: + if state.get("version") != 1: + raise ValueError("Unsupported DeepmdData progress version") + if state.get("identity") != self._progress_identity(): + raise ValueError("DeepmdData progress does not match the data source") + for key in ("set_count", "iterator", "set_size"): + if type(state.get(key)) is not int or state[key] < 0: + raise ValueError(f"Invalid DeepmdData progress {key}") + shuffle_state = state.get("shuffle_state") + if state["set_count"] == 0: + if state["iterator"] or state["set_size"] or shuffle_state is not None: + raise ValueError("Invalid unstarted DeepmdData progress") + elif shuffle_state is None: + raise ValueError("Missing DeepmdData shuffle state") + if shuffle_state is not None: + # Validation uses a private generator, never the shared RNG. + try: + np.random.RandomState(0).set_state(shuffle_state) + except (TypeError, ValueError, IndexError) as error: + raise ValueError("Invalid DeepmdData shuffle state") from error + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore logical progress without loading or scanning frame data. + + A same-set canonical cache may be reused when its requirements have + not changed. New readers reconstruct arrays only on the next batch. + """ + self._validate_state_dict(state) + active_set = ( + self.dirs[(state["set_count"] - 1) % self.get_numb_set()] + if state["set_count"] + else None + ) + keep_set = hasattr(self, "batch_set") and self._batch_set_name == active_set + keep_order = ( + keep_set + and self._batch_shuffle_state == state["shuffle_state"] + and self._batch_set_size == state["set_size"] + ) + if not keep_set: + self.__dict__.pop("batch_set", None) + if not keep_order: + self.__dict__.pop("_batch_indices", None) + self.set_count = state["set_count"] + self.iterator = state["iterator"] + self._batch_set_size = state["set_size"] + self._batch_shuffle_state = copy.deepcopy(state["shuffle_state"]) + def _derive_pair_margin( self, coord: np.ndarray, @@ -493,16 +618,19 @@ def iter_test( chunk_atoms: int, numb_test: float = float("inf"), ) -> Iterator[dict]: - """Yield the test data in chunks of at most ``chunk_atoms`` atoms. + """Yield test data in chunks bounded by ``chunk_atoms`` work units. The complete test system is materialized first, so chunking here bounds the evaluation batch and its predictions rather than source-data memory. + For ordinary systems one frame costs ``natoms`` work units. For + frame-major data such as density grids, the largest point extent also + contributes so grids with many more points than atoms remain bounded. Parameters ---------- chunk_atoms : int - Upper bound on the number of atoms per chunk. A chunk always - carries at least one frame, however many atoms it has. + Upper bound on the per-chunk work proxy. A chunk always carries at + least one frame, however many atoms or frame-major points it has. numb_test : float, optional Upper bound on the number of frames served. A non-finite bound serves the whole test set. @@ -517,7 +645,17 @@ def iter_test( total = int(self.test_set["type"].shape[0]) if np.isfinite(numb_test): total = min(total, int(numb_test)) - step = max(1, int(chunk_atoms) // max(1, self.natoms)) + frame_size = self.natoms + for key, requirement in self.data_dict.items(): + if requirement.get("special_shape") != "frame_major": + continue + value = self.test_set.get(key) + if not isinstance(value, np.ndarray) or value.ndim < 2: + continue + ndof = int(requirement["ndof"]) + extent = int(np.prod(value.shape[1:], dtype=np.int64)) // ndof + frame_size = max(frame_size, extent) + step = max(1, int(chunk_atoms) // max(1, frame_size)) for begin in range(0, total, step): idx = np.arange(begin, min(begin + step, total), dtype=np.int64) chunk = self._get_subdata(self.test_set, idx=idx) @@ -674,12 +812,36 @@ def get_single_frame(self, index: int, num_worker: int) -> dict: # 5. Standardize keys frame_data = {kk.replace("atomic", "atom"): vv for kk, vv in frame_data.items()} - # 6. Reshape atomic data to match expected format [natoms, ndof] + # frame-major keys of one frame must agree on the point extent: a + # grid/density mismatch would otherwise broadcast silently (or fail + # with a raw size error) far downstream in the loss. The per-frame + # payload is (npoints, ndof) for a 3-D source array and (npoints*ndof,) + # for a flattened one. + frame_major_extents = {} + for kk, vv in self.data_dict.items(): + if ( + vv.get("special_shape") == "frame_major" + and kk in frame_data + and float(frame_data.get("find_" + kk, 0.0)) == 1.0 + ): + arr = frame_data[kk] + extent = arr.shape[-2] if arr.ndim >= 2 else arr.shape[-1] // vv["ndof"] + frame_major_extents[kk] = extent + if len(set(frame_major_extents.values())) > 1: + raise ValueError( + f"The frame-major data in {set_dir} have inconsistent " + f"extents {frame_major_extents}; grid and density labels " + "must share the same number of points" + ) + + # 6. Reshape non-atomic data to match expected format [ndof]; + # frame-major arrays keep their per-frame (npoints, ndof) layout for kk in self.data_dict.keys(): if ( "find_" not in kk and kk in frame_data and not self.data_dict[kk]["atomic"] + and self.data_dict[kk].get("special_shape") != "frame_major" ): frame_data[kk] = frame_data[kk].reshape(-1) frame_data["atype"] = frame_data["type"] @@ -789,12 +951,32 @@ def _get_subdata( new_data[ii] = dd return new_data - def _load_batch_set(self, set_name: DPPath) -> None: - if not hasattr(self, "batch_set") or self.get_numb_set() > 1: - self.batch_set = self._load_set(set_name) + def _load_batch_set( + self, set_name: DPPath, *, shuffle_state: tuple | None = None + ) -> None: + if not hasattr(self, "batch_set") or self._batch_set_name != set_name: + batch_set = self._load_set(set_name) if self.modifier is not None: - self.modifier.modify_data(self.batch_set, self) - self.batch_set, _ = self._shuffle_data(self.batch_set) + self.modifier.modify_data(batch_set, self) + self.batch_set = batch_set + self._batch_set_name = set_name + # Always index the canonical source order. Shuffling an already + # shuffled cache cannot be reconstructed from just the latest RNG + # state, and repeatedly expands frames with numb_copy > 1. + nframes = self.batch_set["coord"].shape[0] + self._batch_indices = np.repeat( + np.arange(nframes, dtype=np.int64), + self.batch_set["numb_copy"].reshape(nframes), + ) + if shuffle_state is None: + self._batch_shuffle_state = dp_random.get_state() + dp_random.shuffle(self._batch_indices) + else: + self._batch_shuffle_state = copy.deepcopy(shuffle_state) + rng = np.random.RandomState(0) + rng.set_state(shuffle_state) + rng.shuffle(self._batch_indices) + self._batch_set_size = len(self._batch_indices) self.reset_get_batch() def reset_get_batch(self) -> None: @@ -812,6 +994,26 @@ def _load_test_set(self, shuffle_test: bool) -> None: if "find_" in kk: self.test_set[kk] = test_sets[0][kk] else: + vv = self.data_dict.get(kk, {}) + if vv.get("special_shape") == "frame_major": + # sets of one system must agree on the frame-major extent: + # numpy would otherwise fail with a raw shape error + extents = { + ( + ss[kk].shape[-2] + if ss[kk].ndim >= 3 + else ss[kk].shape[-1] // vv["ndof"] + ) + for ss in test_sets + if float(ss.get("find_" + kk, 0.0)) == 1.0 + } + if len(extents) > 1: + raise ValueError( + f"The frame-major data {kk} has inconsistent " + f"extents {sorted(extents)} across the sets of " + "this system; every set must carry the same " + "number of points" + ) self.test_set[kk] = np.concatenate( [test_set[kk] for test_set in test_sets], axis=0 ) @@ -828,9 +1030,14 @@ def _shuffle_data(self, data: dict[str, Any]) -> dict[str, Any]: for kk in data: if ( isinstance(data[kk], np.ndarray) - and len(data[kk].shape) == 2 and data[kk].shape[0] == nframes and "find_" not in kk + # frame-major arrays (e.g. grid/density) keep their 3D layout + # (nframes, ngrid, ...) and shuffle along the frame axis + and ( + len(data[kk].shape) == 2 + or self.data_dict.get(kk, {}).get("special_shape") == "frame_major" + ) ): ret[kk] = data[kk][idx] else: @@ -923,6 +1130,28 @@ def _load_set(self, set_name: DPPath) -> dict[str, Any]: np.reshape(tmp_in, [nframes, self.natoms, ndof]), axis=1 ) + # frame-major keys that share the same extent must agree on it: + # a grid/density pair with different npoints would only fail far + # downstream when the residual is formed + frame_major_extents = {} + for kk, vv in self.data_dict.items(): + if ( + vv.get("special_shape") == "frame_major" + and kk in data + and float(data.get("find_" + kk, 0.0)) == 1.0 + ): + arr = data[kk] + extent = arr.shape[-2] if arr.ndim >= 3 else arr.shape[-1] // vv["ndof"] + frame_major_extents[kk] = extent + if frame_major_extents: + extents = set(frame_major_extents.values()) + if len(extents) > 1: + raise ValueError( + f"The frame-major data in {set_name} have inconsistent " + f"extents {frame_major_extents}; grid and density labels " + "must share the same number of points" + ) + if self.mixed_type: # nframes x natoms atom_type_mix = self._load_type_mix(set_name) @@ -991,6 +1220,50 @@ def _load_data( special_shape: str | None = None, ) -> np.ndarray: is_hessian = special_shape == "hessian" or key == "hessian" + if dtype is not None: + pass + elif high_prec: + dtype = GLOBAL_ENER_FLOAT_PRECISION + else: + dtype = GLOBAL_NP_FLOAT_PRECISION + path = self._get_data_path(set_name, key) + # Frame-major data does not have an atom axis, so validate and return + # it before preparing atom-selection bookkeeping. + if special_shape == "frame_major" and path.is_file(): + data = path.load_numpy().astype(dtype) + if data.ndim not in (2, 3): + raise ValueError( + f"The frame-major data {key} in {set_name} must have shape " + f"(nframes, npoints * ndof) or (nframes, npoints, ndof), " + f"but got {data.shape}" + ) + if data.shape[0] != nframes: + raise ValueError( + f"The frame count of data {key} in {set_name} is " + f"{data.shape[0]}, which doesn't match the set's nframes {nframes}" + ) + if data.ndim >= 3: + # (nframes, npoints, ndof): trailing dim must be exactly ndof + if data.shape[-1] != ndof_: + raise ValueError( + f"The data {key} in {set_name} has trailing dimension " + f"{data.shape[-1]}, which doesn't match the declared " + f"ndof {ndof_}" + ) + elif data.shape[-1] % ndof_ != 0: + # (nframes, npoints*ndof) flattened: width must be a multiple + raise ValueError( + f"The data {key} in {set_name} has trailing dimension " + f"{data.shape[-1]}, which is not a multiple of the " + f"declared ndof {ndof_}" + ) + npoints = data.shape[1] if data.ndim == 3 else data.shape[1] // ndof_ + if npoints == 0: + raise ValueError( + f"The frame-major data {key} in {set_name} must contain " + "at least one point per frame" + ) + return np.float32(1.0), data if atomic: natoms = self.natoms idx_map = self.idx_map @@ -1008,13 +1281,6 @@ def _load_data( ndof = ndof_ natoms_sel = 0 idx_map_sel = None - if dtype is not None: - pass - elif high_prec: - dtype = GLOBAL_ENER_FLOAT_PRECISION - else: - dtype = GLOBAL_NP_FLOAT_PRECISION - path = self._get_data_path(set_name, key) if path.is_file(): data = path.load_numpy().astype(dtype) try: # YWolfeee: deal with data shape error @@ -1081,7 +1347,12 @@ def _load_data( ndof = 9 * self.natoms * self.natoms elif atomic and type_sel is not None and not output_natoms_for_type_sel: ndof = ndof_ * natoms_sel - data = np.full([nframes, ndof], default, dtype=dtype) + if special_shape == "frame_major": + # frame extent is unknown without the file; keep an empty + # per-frame payload and rely on find=0 downstream + data = np.full([nframes, 0, ndof], default, dtype=dtype) + else: + data = np.full([nframes, ndof], default, dtype=dtype) if repeat != 1: data = np.repeat(data, repeat).reshape([nframes, -1]) return np.float32(0.0), data @@ -1146,6 +1417,10 @@ def _load_single_data( # Create a default array based on requirements if is_hessian: data = np.full([9 * natoms * natoms], vv["default"], dtype=dtype) + elif vv.get("special_shape") == "frame_major": + # frame extent is unknown without the file; an empty frame + # dimension is the honest shape and find=0 marks it absent + data = np.full([0, ndof], vv["default"], dtype=dtype) elif vv["atomic"]: if vv["type_sel"] is not None and not vv["output_natoms_for_type_sel"]: natoms = natoms_sel @@ -1163,13 +1438,53 @@ def _load_single_data( # For filesystem paths, use memmap for better performance mmap_obj = self._get_memmap(path) - # corner case: single frame - if set_nframes == 1: + # Validate frame-major layout before indexing, so that shuffling + # cannot pair a structure with another frame's grid or density label. + if vv.get("special_shape") == "frame_major": + if mmap_obj.ndim not in (2, 3): + raise ValueError( + f"The frame-major data {key} in {set_dir} must have shape " + f"(nframes, npoints * ndof) or (nframes, npoints, ndof), " + f"but got {mmap_obj.shape}" + ) + if mmap_obj.shape[0] != set_nframes: + raise ValueError( + f"The frame count of data {key} in {set_dir} is " + f"{mmap_obj.shape[0]}, which doesn't match the set's " + f"nframes {set_nframes}" + ) + if mmap_obj.ndim >= 3: + if mmap_obj.shape[-1] != ndof: + raise ValueError( + f"The data {key} in {set_dir} has trailing dimension " + f"{mmap_obj.shape[-1]}, which doesn't match the declared " + f"ndof {ndof}" + ) + elif mmap_obj.shape[-1] % ndof != 0: + raise ValueError( + f"The data {key} in {set_dir} has trailing dimension " + f"{mmap_obj.shape[-1]}, which is not a multiple of the " + f"declared ndof {ndof}" + ) + npoints = ( + mmap_obj.shape[1] if mmap_obj.ndim == 3 else mmap_obj.shape[1] // ndof + ) + if npoints == 0: + raise ValueError( + f"The frame-major data {key} in {set_dir} must contain " + "at least one point per frame" + ) + + # corner case: single frame. frame-major arrays always carry a + # leading frame dimension (validated above), so no expansion is needed. + if set_nframes == 1 and vv.get("special_shape") != "frame_major": mmap_obj = mmap_obj[None, ...] # Slice the single frame and make an in-memory copy for modification data = mmap_obj[frame_idx].copy().astype(dtype, copy=False) try: + if vv.get("special_shape") == "frame_major": + return np.float32(1.0), data if is_hessian: data = data.reshape(3 * natoms, 3 * natoms) num_chunks, chunk_size = len(idx_map), 3 @@ -1323,6 +1638,9 @@ class DataRequirementItem: special_shape : str, optional Name of a loader-defined non-standard shape contract. ``"hessian"`` stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame. + ``"frame_major"`` stores one variable-length ``(npoints, ndof)`` array + per frame (e.g. grid/density), kept frame-major without natoms + reshaping. source_policy : {"tracked", "default", "derived"}, optional How source availability affects the consumer. ``"tracked"`` keeps source presence in the ``find_*`` contract. ``"default"`` treats the diff --git a/deepmd/utils/data_config.py b/deepmd/utils/data_config.py new file mode 100644 index 0000000000..917e51849e --- /dev/null +++ b/deepmd/utils/data_config.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Format-independent configuration of training data groups.""" + +import logging +from collections.abc import ( + Mapping, + Sequence, +) +from dataclasses import ( + dataclass, +) +from typing import ( + Any, + TypeAlias, +) + +from deepmd.utils.data_path import ( + DatasetFormat, + is_lmdb, + process_systems, +) +from deepmd.utils.data_sampling import ( + group_probabilities, +) +from deepmd.utils.path import ( + DPPath, +) + +DataConfig: TypeAlias = Mapping[str, Any] | list[Mapping[str, Any]] +BatchSize: TypeAlias = int | str | list[int] +log = logging.getLogger(__name__) + + +def check_type_map_consistency(type_maps: Sequence[list[str] | None]) -> list[str]: + """Return the longest mutually consistent type map.""" + result: list[str] = [] + for type_map in type_maps: + if type_map is None: + continue + common = min(len(result), len(type_map)) + if result[:common] != type_map[:common]: + raise RuntimeError( + f"Inconsistent type map: {result!s} {type_map!s} in different data systems. " + "If you didn't set model/type_map, please set it, " + "since the type map of the model cannot be decided by data." + ) + if len(type_map) > len(result): + result = list(type_map) + return result + + +def data_group_configs(config: DataConfig) -> list[Mapping[str, Any]]: + """View a legacy configuration as one group without mutating its input.""" + groups = [config] if isinstance(config, Mapping) else config + if not isinstance(groups, list) or not groups: + raise ValueError( + "Data configuration must be a dictionary or a nonempty list of dictionaries." + ) + if any(not isinstance(group, Mapping) for group in groups): + raise ValueError("Every data group must be a dictionary.") + group_probabilities([1] * len(groups), [group.get("prob") for group in groups]) + return groups + + +@dataclass(frozen=True) +class DataGroupConfig: + """One configured group after recursive dataset discovery. + + ``prob`` controls the whole group's batch-selection probability. + ``auto_prob`` and ``sys_probs`` retain their native within-group + semantics. An LMDB's internal system numbering is local to that file; + custom ``auto_prob`` rules therefore require a single LMDB source. + """ + + systems: tuple[str, ...] + batch_size: BatchSize + auto_prob: str | None + sys_probs: list[float] | None + prob: float | None + + +@dataclass(frozen=True) +class DataSourceConfig: + """One native batch stream within a configured group.""" + + format: DatasetFormat + systems: tuple[str, ...] + batch_size: BatchSize + + +def resolve_data_sources(group: DataGroupConfig) -> list[DataSourceConfig]: + """Partition a group into native readers after whole-system filtering. + + NumPy systems share a native reader; each LMDB owns one. NumPy atom-count + filtering precedes this partition so an empty NumPy subset does not reject + a group with retained LMDB data. LMDB filtering remains frame-local. + """ + numpy_indices = [ + index for index, path in enumerate(group.systems) if not is_lmdb(path) + ] + numpy_batch_size = group.batch_size + if isinstance(numpy_batch_size, str) and numpy_batch_size.startswith("filter:"): + limit = int(numpy_batch_size.split(":")[1]) + retained = [ + index + for index in numpy_indices + if (DPPath(group.systems[index]) / "type.raw").load_txt(ndmin=1).size + <= limit + ] + if len(retained) != len(numpy_indices): + log.warning( + "Remove %d systems with more than %d atoms", + len(numpy_indices) - len(retained), + limit, + ) + numpy_indices = retained + numpy_batch_size = f"max:{limit}" + elif isinstance(numpy_batch_size, list): + numpy_batch_size = [numpy_batch_size[index] for index in numpy_indices] + sources = [] + if numpy_indices: + sources.append( + DataSourceConfig( + "deepmd", + tuple(group.systems[index] for index in numpy_indices), + numpy_batch_size, + ) + ) + for index, path in enumerate(group.systems): + if is_lmdb(path): + batch_size = group.batch_size + if isinstance(batch_size, list): + batch_size = batch_size[index] + sources.append(DataSourceConfig("lmdb", (path,), batch_size)) + if not sources: + raise ValueError( + f"No system left in data group after applying {group.batch_size!r}." + ) + return sources + + +def resolve_data_groups( + config: DataConfig, + *, + formats: tuple[DatasetFormat, ...] = ("deepmd", "lmdb"), +) -> list[DataGroupConfig]: + """Resolve every group's paths and validate format-dependent options.""" + result = [] + for index, group in enumerate(data_group_configs(config)): + roots = ( + [group["systems"]] + if isinstance(group["systems"], str) + else group["systems"] + ) + if "lmdb" not in formats and any(is_lmdb(path) for path in roots): + raise ValueError("LMDB training data requires the pt or pt_expt backend.") + systems = process_systems( + group["systems"], group.get("rglob_patterns"), formats=formats + ) + if not systems: + raise ValueError(f"Data group {index} contains no supported data systems.") + if len(systems) > 1 and any(is_lmdb(path) for path in systems): + if group.get("auto_prob", "prob_sys_size") not in ( + "prob_sys_size", + "prob_uniform", + ): + raise ValueError( + f"Data group {index} expands to multiple sources including LMDB, " + "so auto_prob must be prob_sys_size or prob_uniform. " + "Use separate training_data groups " + "to configure each LMDB's internal system probabilities." + ) + if group.get("sys_probs") is not None: + raise ValueError( + f"Data group {index} expands to multiple sources including LMDB, " + "so sys_probs is not supported. Use training_data groups with prob " + "to assign source probabilities." + ) + batch_size = group.get("batch_size", "auto") + if isinstance(batch_size, list) and len(batch_size) != len(systems): + raise ValueError( + f"Data group {index} has {len(systems)} expanded systems but " + f"{len(batch_size)} batch_size entries." + ) + result.append( + DataGroupConfig( + systems=tuple(systems), + batch_size=batch_size, + auto_prob=group.get("auto_prob"), + sys_probs=group.get("sys_probs"), + prob=group.get("prob"), + ) + ) + return result diff --git a/deepmd/utils/data_group.py b/deepmd/utils/data_group.py new file mode 100644 index 0000000000..df9b15b0bc --- /dev/null +++ b/deepmd/utils/data_group.py @@ -0,0 +1,626 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Batch-level composition of independently configured data systems.""" + +from __future__ import ( + annotations, +) + +import logging +from copy import ( + deepcopy, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +import numpy as np + +from deepmd.dpmodel.utils.training_utils import ( + compute_total_numb_batch, + get_training_data_size, +) +from deepmd.utils.data_config import ( + check_type_map_consistency, +) +from deepmd.utils.data_sampling import ( + data_child_seeds, + group_probabilities, +) +from deepmd.utils.model_stat import ( + _get_stat_batch, + _get_stat_nsystems, + _get_stat_pass_length, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + Iterator, + Mapping, + Sequence, + ) + + from deepmd.utils.data import ( + DataRequirementItem, + ) + +log = logging.getLogger(__name__) + + +class DataBatch(dict[str, Any]): + """Carry source-specific execution context without adding model input fields.""" + + def __init__(self, values: Mapping[str, Any], pair_filter_enabled: bool) -> None: + super().__init__(values) + self.pair_filter_enabled = pair_filter_enabled + + +def iter_data_sources(data: Any) -> Iterator[Any]: + """Visit each native data source of a composed dataset in order.""" + children = getattr(data, "children", None) + if children is not None: + for child in children: + yield from iter_data_sources(child) + else: + yield data + + +def epoch_batch_counts(data: Any, world_size: int) -> np.ndarray: + """Resolve local workloads while preserving the replicated-reader contract.""" + get_counts = getattr(data, "get_epoch_batch_counts", None) + if get_counts is not None: + return get_counts(world_size) + return np.asarray(data.nbatches, dtype=np.float64) / world_size + + +class GroupedDataSystem: + """Select whole batches while preserving every child's batching rules. + + The child owns its within-system probabilities, batching, and I/O. This + node owns only the probability of selecting that child. Training and + statistical system indices are separate flattened views: an LMDB is one + training source but may expose several homogeneous statistical systems. + + Uniform composition selects each native source with weight ``M / n``, + where ``M`` counts its retained systems and ``n`` is its next batch's + frame count, or its mean batch size for a persistent fixed-pass stream. + Each system therefore supplies the same long-run frame share despite + different native batch sizes. Nominal epoch weights use pass means; + runtime weights can change as LMDB packing advances. + """ + + def __init__( + self, + children: Sequence[Any], + probabilities: Sequence[float | None], + *, + seed: int | None = None, + uniform: bool = False, + ) -> None: + if not children or len(children) != len(probabilities): + raise ValueError("Data groups require one probability per nonempty child.") + self.children = tuple(children) + self._assigned_probabilities = tuple(probabilities) + self._uniform = uniform + self._rng = np.random.default_rng(data_child_seeds(seed, 1)[0]) + self._probabilities: np.ndarray | None = None + self._cdf: np.ndarray | None = None + self._filter_enabled: tuple[bool, ...] | None = None + self._remaining = 0 + self._pass_length = 0 + self._align_metadata() + + def _align_metadata(self) -> None: + """Resolve metadata shared by the native data sources.""" + sources = list(iter_data_sources(self)) + self.type_map = check_type_map_consistency( + [source.get_type_map() for source in sources] + ) + self._ntypes = max(source.get_ntypes() for source in sources) + self.mixed_type = any(source.mixed_type for source in sources) + + @property + def component_probabilities(self) -> np.ndarray: + """Return this sampling pass's probabilities over child datasets.""" + if self._probabilities is None: + self._probabilities = group_probabilities( + [ + child.uniform_nominal_mass if self._uniform else sum(child.nbatches) + for child in self.children + ], + self._assigned_probabilities, + ) + return self._probabilities + + @property + def nbatches(self) -> list[int]: + """Return global batch counts in the flattened training-system view.""" + return [count for child in self.children for count in child.nbatches] + + @property + def sys_probs(self) -> np.ndarray: + """Return global probabilities without duplicating child sampling state.""" + return np.concatenate( + [ + probability * np.asarray(child.sys_probs) + for child, probability in zip( + self.children, self.component_probabilities, strict=True + ) + ] + ) + + def set_filter_enabled(self, enabled: Sequence[bool]) -> None: + """Bind the configured frame-filter policy to each child stream.""" + if len(enabled) != len(self.children): + raise ValueError("Each data group requires one frame-filter policy.") + self._filter_enabled = tuple(enabled) + + def _batch_context(self, batch: dict[str, Any], child: int) -> dict[str, Any]: + if self._filter_enabled is None: + return batch + return DataBatch(batch, self._filter_enabled[child]) + + def _invalidate_sampling(self) -> None: + self._probabilities = None + self._cdf = None + self._remaining = 0 + + def _start_pass(self) -> None: + self._probabilities = None + probabilities = self.component_probabilities + self._cdf = np.cumsum(probabilities) + self._cdf[-1] = 1.0 + self._pass_length = self._reader_pass_length() + self._remaining = self._pass_length + + @staticmethod + def _locate(index: int, sizes: Sequence[int]) -> tuple[int, int]: + if index < 0: + raise IndexError(f"System index {index} must be non-negative.") + for child, size in enumerate(sizes): + if index < size: + return child, index + index -= size + raise IndexError("System index is outside the grouped dataset.") + + def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: + """Draw a child batch, or address one flattened training system.""" + if sys_idx is not None: + child, local = self._locate( + sys_idx, [source.get_nsystems() for source in self.children] + ) + return self._batch_context( + self.children[child].get_batch(sys_idx=local), child + ) + if self._remaining == 0: + self._start_pass() + assert self._cdf is not None + cdf = self._cdf + if self._uniform: + probabilities = group_probabilities( + [child.uniform_sampling_mass for child in self.children], + self._assigned_probabilities, + ) + cdf = np.cumsum(probabilities) + cdf[-1] = 1.0 + index = int(np.searchsorted(cdf, self._rng.random(), side="right")) + self._remaining -= 1 + return self._batch_context(self.children[index].get_batch(), index) + + def _sampling_config(self) -> dict[str, Any]: + return { + "probabilities": [ + None if value is None else float(value) + for value in self._assigned_probabilities + ], + "uniform": self._uniform, + } + + def state_dict(self) -> dict[str, Any]: + """Save this node's selection stream and each child's logical cursor.""" + return { + "version": 1, + "config": self._sampling_config(), + "rng": deepcopy(self._rng.bit_generator.state), + "remaining": self._remaining, + "pass_length": self._pass_length, + "probabilities": ( + self.component_probabilities.tolist() if self._remaining else None + ), + "children": [child.state_dict() for child in self.children], + } + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore the hierarchy without drawing batches or decoding frames.""" + if state.get("version") != 1: + raise ValueError("Unsupported grouped-data progress version") + if state.get("config") != self._sampling_config(): + raise ValueError("Grouped data progress has different sampling settings") + children = state.get("children") + if not isinstance(children, list) or len(children) != len(self.children): + raise ValueError("Grouped data progress has different children") + for key in ("remaining", "pass_length"): + if type(state.get(key)) is not int or state[key] < 0: + raise ValueError(f"Grouped data progress {key} must be non-negative") + if state["remaining"] > state["pass_length"]: + raise ValueError("Grouped data progress exceeds its pass length") + probabilities = state.get("probabilities") + if probabilities is not None: + probabilities = np.asarray(probabilities, dtype=np.float64) + if ( + probabilities.shape != (len(self.children),) + or not np.all(np.isfinite(probabilities)) + or np.any(probabilities < 0) + or not np.isclose(probabilities.sum(), 1.0, rtol=1e-12, atol=0) + ): + raise ValueError("Grouped data progress has invalid probabilities") + elif state["remaining"]: + raise ValueError( + "Grouped data progress is missing active-pass probabilities" + ) + if "rng" not in state: + raise ValueError("Grouped data progress is missing its selection RNG") + rng = np.random.default_rng(0) + rng.bit_generator.state = state["rng"] + for child, progress in zip(self.children, children, strict=True): + child.load_state_dict(progress) + self._rng = rng + self._remaining = state["remaining"] + self._pass_length = state["pass_length"] + self._probabilities = probabilities + self._cdf = None if probabilities is None else np.cumsum(probabilities) + if self._cdf is not None: + self._cdf[-1] = 1.0 + + def _reader_pass_length(self) -> int: + return compute_total_numb_batch( + [child.get_batch_pass_length() for child in self.children], + self.component_probabilities, + ) + + def get_batch_pass_length(self) -> int: + """Return the local reader-pass length belonging to the current draw.""" + return self._pass_length or self._reader_pass_length() + + def get_epoch_batch_counts(self, world_size: int) -> np.ndarray: + """Expose local optimizer workloads, including each leaf's sharding.""" + return np.concatenate( + [epoch_batch_counts(child, world_size) for child in self.children] + ) + + def get_stat_nsystems(self) -> int: + return sum(_get_stat_nsystems(child) for child in self.children) + + def get_stat_numb_batches(self, sys_idx: int) -> int: + child, local = self._locate( + sys_idx, [_get_stat_nsystems(source) for source in self.children] + ) + return _get_stat_pass_length(self.children[child], local) + + def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: + child, local = self._locate( + sys_idx, [_get_stat_nsystems(source) for source in self.children] + ) + return self._batch_context(_get_stat_batch(self.children[child], local), child) + + def add_data_requirements(self, requirements: list[DataRequirementItem]) -> None: + for child in self.children: + child.add_data_requirements(requirements) + self._invalidate_sampling() + + def use_ragged_batches(self, ragged: bool) -> None: + for child in self.children: + configure = getattr(child, "use_ragged_batches", None) + if configure is not None: + configure(ragged) + self._invalidate_sampling() + + def set_frame_transform(self, transform: Any) -> None: + for child in self.children: + configure = getattr(child, "set_frame_transform", None) + if configure is None: + raise ValueError("Frame transforms require LMDB data in every group.") + configure(transform) + + def close(self) -> None: + for child in self.children: + close = getattr(child, "close", None) + if close is not None: + close() + + def get_nsystems(self) -> int: + return sum(child.get_nsystems() for child in self.children) + + def get_model_sampling_size(self) -> int: + """Sum native data sizes without applying within-task batch probabilities.""" + return sum(get_training_data_size(child) for child in self.children) + + def get_ntypes(self) -> int: + return self._ntypes + + def get_type_map(self) -> list[str]: + return self.type_map + + def get_nbatches(self) -> list[int]: + return self.nbatches + + def get_batch_size(self) -> list[int]: + return [size for child in self.children for size in child.get_batch_size()] + + def get_sys(self, index: int) -> Any: + child, local = self._locate( + index, [source.get_nsystems() for source in self.children] + ) + return self.children[child].get_sys(local) + + def get_data_dict(self, ii: int = 0) -> dict[str, Any]: + child, local = self._locate( + ii, [source.get_nsystems() for source in self.children] + ) + return self.children[child].get_data_dict(local) + + @property + def data_systems(self) -> list[Any]: + return [system for child in self.children for system in child.data_systems] + + @property + def system_dirs(self) -> list[str]: + return [path for child in self.children for path in child.system_dirs] + + def add_dict(self, requirements: dict[str, dict[str, Any]]) -> None: + for child in self.children: + child.add_dict(requirements) + + def add(self, *args: Any, **kwargs: Any) -> None: + for child in self.children: + child.add(*args, **kwargs) + + def reduce(self, key_out: str, key_in: str) -> None: + for child in self.children: + child.reduce(key_out, key_in) + + def print_summary(self, name: str, prob: Any = None) -> None: + for index, (child, probability) in enumerate( + zip(self.children, self.component_probabilities, strict=True) + ): + label = f"{name}/group_{index}" + log.info("%s: batch probability %.6g", label, probability) + child.print_summary(label) + + +class GroupedDataset: + """Own native loader datasets and expose their combined statistics view.""" + + def __init__( + self, + children: Sequence[Any], + configs: Sequence[Mapping[str, Any]], + probabilities: Sequence[float | None], + *, + uniform: bool = False, + ) -> None: + if not children or not len(children) == len(configs) == len(probabilities): + raise ValueError( + "A dataset group requires matching datasets, configs, and probabilities." + ) + self.children = tuple(children) + self.configs = tuple(configs) + self.probability_specs = tuple(probabilities) + self.uniform = uniform + self.filter_enabled: tuple[bool, ...] | None = None + + def __len__(self) -> int: + """Preserve native dataset sizes for default multi-task probabilities.""" + return sum(len(child) for child in self.children) + + def set_filter_enabled(self, enabled: Sequence[bool]) -> None: + """Bind the configured frame-filter policy to each child stream.""" + if len(enabled) != len(self.children): + raise ValueError("Each data group requires one frame-filter policy.") + self.filter_enabled = tuple(enabled) + + @property + def stat_filter_enabled(self) -> list[bool] | None: + if self.filter_enabled is None: + return None + return [ + enabled + for child, enabled in zip(self.children, self.filter_enabled, strict=True) + for _ in child.systems + ] + + @property + def systems(self) -> list[Any]: + return [system for child in self.children for system in child.systems] + + @property + def dataloaders(self) -> list[Any]: + return [loader for child in self.children for loader in child.dataloaders] + + @property + def sampler_list(self) -> list[Any]: + return [sampler for child in self.children for sampler in child.sampler_list] + + @property + def index(self) -> list[int]: + return [count for child in self.children for count in child.index] + + @property + def total_batch(self) -> int: + return sum(self.index) + + def add_data_requirement(self, requirements: list[DataRequirementItem]) -> None: + for child in self.children: + child.add_data_requirement(requirements) + + def close(self) -> None: + for child in self.children: + close = getattr(child, "close", None) + if close is not None: + close() + + +class TrainingDataLoader: + """Expose native loader iteration, sampling metadata, and resource ownership.""" + + def __init__( + self, + dataset: Any, + loader: Any, + probabilities: np.ndarray, + batch_counts: Callable[[], Sequence[int]], + *, + sampling_sizes: Callable[[], Sequence[int]] | None = None, + uniform_sampling: Any = None, + ) -> None: + self.dataset = dataset + self.loader = loader + self.probabilities = np.asarray(probabilities, dtype=np.float64) + self._batch_counts = batch_counts + self._sampling_sizes = sampling_sizes or batch_counts + self._uniform_sampling = uniform_sampling + + @property + def uniform_sampling_mass(self) -> float: + return self._uniform_sampling.uniform_sampling_mass + + @property + def uniform_nominal_mass(self) -> float: + return self._uniform_sampling.uniform_nominal_mass + + @property + def batch_counts(self) -> np.ndarray: + return np.asarray(self._batch_counts(), dtype=np.float64) + + @property + def sampling_sizes(self) -> np.ndarray: + """Return natural source masses, independently of per-rank padding.""" + return np.asarray(self._sampling_sizes(), dtype=np.float64) + + @property + def system_probabilities(self) -> np.ndarray: + return self.probabilities + + @property + def epoch_length(self) -> int: + return compute_total_numb_batch(self.batch_counts, self.system_probabilities) + + def __iter__(self) -> Iterator[dict[str, Any]]: + return iter(self.loader) + + def __len__(self) -> int: + return len(self.loader) + + def print_summary(self, name: str) -> None: + self.dataset.print_summary(name, self.probabilities) + + def close(self) -> None: + close_loader = getattr(self.loader, "close", None) + if close_loader is not None: + close_loader() + close_dataset = getattr(self.dataset, "close", None) + if close_dataset is not None: + close_dataset() + + +class GroupedTrainingDataLoader: + """Interleave native batch streams without moving decoding into a wrapper worker.""" + + def __init__( + self, + dataset: GroupedDataset, + children: Sequence[TrainingDataLoader | GroupedTrainingDataLoader], + seed: int | None, + ) -> None: + self.dataset = dataset + self.children = tuple(children) + self._rng = np.random.default_rng(seed) + self._iterators: dict[int, Iterator[dict[str, Any]]] = {} + + @property + def probabilities(self) -> np.ndarray: + return group_probabilities( + [ + child.uniform_nominal_mass + if self.dataset.uniform + else float(child.sampling_sizes.sum()) + for child in self.children + ], + self.dataset.probability_specs, + ) + + @property + def batch_counts(self) -> np.ndarray: + return np.concatenate([child.batch_counts for child in self.children]) + + @property + def sampling_sizes(self) -> np.ndarray: + return np.concatenate([child.sampling_sizes for child in self.children]) + + @property + def system_probabilities(self) -> np.ndarray: + return np.concatenate( + [ + probability * child.system_probabilities + for child, probability in zip( + self.children, self.probabilities, strict=True + ) + ] + ) + + @property + def epoch_length(self) -> int: + return compute_total_numb_batch(self.batch_counts, self.system_probabilities) + + def __len__(self) -> int: + return compute_total_numb_batch( + [len(child) for child in self.children], self.probabilities + ) + + def __iter__(self) -> Iterator[dict[str, Any]]: + probabilities = self.probabilities + cdf = np.cumsum(probabilities) + cdf[-1] = 1.0 + count = compute_total_numb_batch( + [len(child) for child in self.children], probabilities + ) + offsets = np.cumsum([0, *(len(child.batch_counts) for child in self.children)]) + for _ in range(count): + if self.dataset.uniform: + probabilities = group_probabilities( + [child.uniform_sampling_mass for child in self.children], + self.dataset.probability_specs, + ) + cdf = np.cumsum(probabilities) + cdf[-1] = 1.0 + index = int(np.searchsorted(cdf, self._rng.random(), side="right")) + if index not in self._iterators: + self._iterators[index] = iter(self.children[index]) + batch = next(self._iterators[index], None) + if batch is None: + self._iterators[index] = iter(self.children[index]) + batch = next(self._iterators[index]) + result = {**batch, "sid": batch["sid"] + int(offsets[index])} + if self.dataset.filter_enabled is not None: + result = DataBatch(result, self.dataset.filter_enabled[index]) + elif isinstance(batch, DataBatch): + result = DataBatch(result, batch.pair_filter_enabled) + yield result + + def print_summary(self, name: str) -> None: + for index, (child, probability) in enumerate( + zip(self.children, self.probabilities, strict=True) + ): + label = f"{name}/group_{index}" + log.info("%s: batch probability %.6g", label, probability) + child.print_summary(label) + + def close(self) -> None: + for iterator in self._iterators.values(): + close = getattr(iterator, "close", None) + if close is not None: + close() + self._iterators.clear() + for child in self.children: + child.close() diff --git a/deepmd/utils/data_path.py b/deepmd/utils/data_path.py new file mode 100644 index 0000000000..4559121a3f --- /dev/null +++ b/deepmd/utils/data_path.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-independent discovery of on-disk data systems.""" + +from pathlib import ( + Path, +) +from typing import ( + Literal, +) + +from deepmd.utils.path import ( + DPOSPath, + DPPath, +) + +DatasetFormat = Literal["deepmd", "lmdb"] + + +def is_lmdb(systems: str) -> bool: + """Return whether a path identifies an LMDB dataset.""" + return systems.endswith(".lmdb") or Path(systems, "data.mdb").is_file() + + +def _inside_lmdb(path: Path, root: Path) -> bool: + for parent in path.parents: + if parent == root: + break + if is_lmdb(str(parent)): + return True + return False + + +def expand_sys_str( + root_dir: str | Path, + *, + formats: tuple[DatasetFormat, ...] = ("deepmd",), + patterns: list[str] | None = None, +) -> list[str]: + """Discover data systems beneath one directory or HDF5 container. + + Parameters + ---------- + root_dir : str or Path + A data system, a parent directory, or a DeePMD HDF5 path. + formats : tuple of str + Accepted leaf formats. ``deepmd`` recognizes ``type.raw`` in NPY + directories and HDF5 groups; ``lmdb`` recognizes LMDB directories. + patterns : list of str, optional + Recursive glob patterns selecting system paths. A directly supplied + LMDB dataset is terminal and does not require a glob match. + + Returns + ------- + list[str] + Discovered system paths. Unpatterned DeePMD traversal retains its + established order, with the root system last. Overlapping patterns + include each matched path once. + """ + root_name = str(root_dir) + if is_lmdb(root_name): + return [root_name] if "lmdb" in formats else [] + + root = DPPath(root_name) + candidates = ( + [*root.rglob("*"), root] + if patterns is None + else [path for pattern in patterns for path in root.rglob(pattern)] + ) + systems: dict[str, None] = {} + for candidate in candidates: + name = str(candidate) + if isinstance(candidate, DPOSPath): + path = candidate.path + if not candidate.is_dir(): + continue + if path != root.path and _inside_lmdb(path, root.path): + continue + if is_lmdb(name): + if "lmdb" in formats: + systems[name] = None + continue + if "deepmd" in formats and (candidate / "type.raw").is_file(): + systems[name] = None + return list(systems) + + +def rglob_sys_str(root_dir: str, patterns: list[str]) -> list[str]: + """Discover DeePMD systems matching recursive glob patterns.""" + return expand_sys_str(root_dir, patterns=patterns) + + +def process_systems( + systems: str | list[str], + patterns: list[str] | None = None, + *, + formats: tuple[DatasetFormat, ...] = ("deepmd", "lmdb"), + root_dir: str | Path | None = None, +) -> list[str]: + """Expand every input path through the same format-aware discovery. + + Explicit input order and repeated input paths are preserved, so an + input list defines the order used by per-system sampling probabilities. + Each entry may name a dataset directly or a parent containing datasets. + When ``root_dir`` is supplied, relative filesystem paths are resolved + against it before discovery; HDF5 group fragments remain unchanged. + """ + paths = [systems] if isinstance(systems, str) else systems + if not isinstance(paths, list) or any(not isinstance(path, str) for path in paths): + raise ValueError( + f"Invalid systems value {systems!r}: expected str or list[str]." + ) + if root_dir is not None: + paths = [ + str((Path(root_dir) / file_name).resolve()) + separator + group + for file_name, separator, group in (path.partition("#") for path in paths) + ] + return [ + system + for path in paths + for system in expand_sys_str(path, formats=formats, patterns=patterns) + ] diff --git a/deepmd/utils/data_sampling.py b/deepmd/utils/data_sampling.py new file mode 100644 index 0000000000..ab22ddb918 --- /dev/null +++ b/deepmd/utils/data_sampling.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Sampling probabilities shared by data formats and training backends.""" + +from collections.abc import ( + Sequence, +) + +import numpy as np + +from deepmd.utils import random as dp_random + + +def process_sys_probs( + sys_probs: Sequence[float], nbatch: Sequence[int] | np.ndarray +) -> np.ndarray: + """Fill unspecified probabilities in proportion to the batch counts. + + Non-negative entries specify absolute probabilities. Negative entries + share the remaining probability in proportion to their batch counts. + A fully assigned probability vector does not depend on batch counts. + """ + probabilities = np.asarray(sys_probs) + if not np.issubdtype(probabilities.dtype, np.floating): + probabilities = probabilities.astype(np.float64) + if probabilities.ndim != 1: + raise ValueError("Sampling probabilities must be one-dimensional.") + if not np.all(np.isfinite(probabilities)): + raise ValueError("Sampling probabilities must be finite.") + assigned = probabilities >= 0 + assigned_sum = float(probabilities[assigned].sum()) + tolerance = max(1e-8, float(np.finfo(probabilities.dtype).eps)) + if assigned_sum > 1.0 + tolerance: + raise ValueError("The sum of assigned probabilities must not exceed 1.") + result = np.where(assigned, probabilities, 0.0) + remaining = max(0.0, 1.0 - assigned_sum) + if remaining > tolerance: + sizes = np.asarray(nbatch, dtype=np.float64) + if probabilities.shape != sizes.shape: + raise ValueError( + "Sampling probabilities and batch counts must have the same length." + ) + if not np.all(np.isfinite(sizes)) or np.any(sizes < 0): + raise ValueError("Batch counts must be finite and non-negative.") + remaining_size = float(sizes[~assigned].sum()) + if remaining_size <= 0: + raise ValueError( + "The remaining probability requires at least one unassigned, " + "nonempty dataset." + ) + result[~assigned] = remaining * sizes[~assigned] / remaining_size + if result.size == 0 or result.sum() <= 0: + raise ValueError("Sampling probabilities must contain a positive entry.") + return result / result.sum() + + +def prob_sys_size_ext( + keywords: str, nsystems: int, nbatch: Sequence[int] | np.ndarray +) -> np.ndarray: + """Distribute each configured block's weight by its system sizes.""" + blocks = [part.split(":") for part in keywords.split(";")[1:]] + weights = np.asarray([float(block[2]) for block in blocks]) + if np.any(weights < 0): + raise ValueError("The weight of a block must be non-negative.") + block_probs = weights / weights.sum() + sys_probs = np.zeros(nsystems, dtype=np.float64) + sizes = np.asarray(nbatch, dtype=np.float64) + for block, probability in zip(blocks, block_probs, strict=True): + start, end = int(block[0]), int(block[1]) + block_sizes = sizes[start:end] + sys_probs[start:end] = block_sizes / block_sizes.sum() * probability + return sys_probs + + +def uniform_sampling_weights(batch_nframes: Sequence[float] | np.ndarray) -> np.ndarray: + """Return inverse batch sizes, excluding streams that supply no frames.""" + sizes = np.asarray(batch_nframes, dtype=np.float64) + return np.divide(1.0, sizes, out=np.zeros_like(sizes), where=sizes > 0) + + +def system_probabilities( + nbatches: Sequence[int] | np.ndarray, + *, + sys_probs: Sequence[float] | None = None, + auto_prob: str = "prob_sys_size", + batch_nframes: Sequence[float] | np.ndarray | None = None, +) -> np.ndarray: + """Resolve batch probabilities, correcting uniform sampling for frame counts. + + ``batch_nframes`` gives the mean frames per batch of each native stream. + Its reciprocal makes ``prob_uniform`` allocate equal sampled frame shares + rather than equal batch counts. Without batching metadata, every draw is + treated as one frame. Explicit ``sys_probs`` remain batch probabilities. + """ + nsystems = len(nbatches) + if nsystems == 0: + raise ValueError("No systems provided.") + if sys_probs is not None: + return process_sys_probs(sys_probs, nbatches) + if auto_prob == "prob_uniform": + sizes = ( + np.ones(nsystems, dtype=np.float64) + if batch_nframes is None + else batch_nframes + ) + weights = uniform_sampling_weights(sizes) + return group_probabilities(weights, [None] * nsystems) + if auto_prob == "prob_sys_size": + auto_prob = f"prob_sys_size;0:{nsystems}:1.0" + if not auto_prob.startswith("prob_sys_size;"): + raise RuntimeError(f"Unknown auto prob style: {auto_prob}") + return prob_sys_size_ext(auto_prob, nsystems, nbatches) + + +def group_probabilities( + sizes: Sequence[float] | np.ndarray, probabilities: Sequence[float | None] +) -> np.ndarray: + """Assign group probabilities, sharing any remainder by group size.""" + if len(sizes) != len(probabilities): + raise ValueError("Group sizes and probabilities must have the same length.") + if any( + probability is not None and probability < 0 for probability in probabilities + ): + raise ValueError("A group's prob must be non-negative or omitted.") + resolved = process_sys_probs( + np.asarray( + [ + -1.0 if probability is None else probability + for probability in probabilities + ], + dtype=np.float64, + ), + sizes, + ) + if np.any((np.asarray(sizes) == 0) & (resolved > 0)): + raise ValueError("A data group with positive probability must contain batches.") + return resolved + + +def data_child_seeds(seed: int | None, count: int) -> list[int]: + """Derive independent selection streams for nodes in a dataset hierarchy.""" + if seed is None: + seed = int(dp_random.choice(2**32)) + return [ + int(child.generate_state(1, dtype=np.uint64)[0]) + for child in np.random.SeedSequence(seed).spawn(count) + ] + + +def global_batch_counts( + frame_counts: Sequence[int], batch_sizes: Sequence[int] +) -> list[int]: + """Count source batches before distributed sampling and rank-local padding.""" + return [ + (frames + batch_size - 1) // batch_size + for frames, batch_size in zip(frame_counts, batch_sizes, strict=True) + ] diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 890224e063..5fbe657150 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -2,9 +2,15 @@ import collections import logging import warnings +from collections.abc import ( + Callable, +) from functools import ( cached_property, ) +from types import ( + SimpleNamespace, +) from typing import ( Any, ) @@ -13,9 +19,7 @@ import deepmd.utils.random as dp_random from deepmd.common import ( - expand_sys_str, make_default_mesh, - rglob_sys_str, ) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, @@ -24,6 +28,25 @@ DataRequirementItem, DeepmdData, ) +from deepmd.utils.data_config import ( + DataConfig, + check_type_map_consistency, + resolve_data_groups, + resolve_data_sources, +) +from deepmd.utils.data_path import ( + is_lmdb, +) +from deepmd.utils.data_path import process_systems as process_systems +from deepmd.utils.data_sampling import ( + data_child_seeds, +) +from deepmd.utils.data_sampling import prob_sys_size_ext as prob_sys_size_ext +from deepmd.utils.data_sampling import process_sys_probs as process_sys_probs +from deepmd.utils.data_sampling import ( + system_probabilities, + uniform_sampling_weights, +) from deepmd.utils.out_stat import ( compute_stats_from_redu, ) @@ -84,7 +107,7 @@ def __init__( automatically by the number of batches in the system. auto_prob_style : str Determine the probability of systems automatically. The method is assigned by this key and can be - - "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems() + - "prob_uniform" : every system supplies the same sampled frame share. - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is divided into blocks. A block is specified by `stt_idx:end_idx:weight`, @@ -230,7 +253,7 @@ def __init__( self.data_systems[ii].get_sys_numb_batch(self.batch_size[ii]) ) type_map_list.append(self.data_systems[ii].get_type_map()) - self.type_map = self._check_type_map_consistency(type_map_list) + self.type_map = check_type_map_consistency(type_map_list) # ! altered by Marián Rynik # test size @@ -434,23 +457,41 @@ def set_sys_probs( sys_probs: list[float] | None = None, auto_prob_style: str = "prob_sys_size", ) -> None: - if sys_probs is None: - if auto_prob_style == "prob_uniform": - prob_v = 1.0 / float(self.nsystems) - probs = [prob_v for ii in range(self.nsystems)] - elif auto_prob_style[:13] == "prob_sys_size": - if auto_prob_style == "prob_sys_size": - prob_style = f"prob_sys_size;0:{self.get_nsystems()}:1.0" - else: - prob_style = auto_prob_style - probs = prob_sys_size_ext( - prob_style, self.get_nsystems(), self.nbatches - ) - else: - raise RuntimeError("Unknown auto prob style: " + auto_prob_style) - else: - probs = process_sys_probs(sys_probs, self.nbatches) - self.sys_probs = probs + self.sys_probs = system_probabilities( + self.nbatches, + sys_probs=sys_probs, + auto_prob=auto_prob_style, + batch_nframes=( + self.mean_batch_nframes + if sys_probs is None and auto_prob_style == "prob_uniform" + else None + ), + ) + + @cached_property + def mean_batch_nframes(self) -> np.ndarray: + """Return native mean batch sizes for per-system selection.""" + if self.mixed_systems: + return np.ones(self.nsystems, dtype=np.float64) + return np.asarray( + [ + system.get_mean_batch_nframes(size) + for system, size in zip(self.data_systems, self.batch_size, strict=True) + ], + dtype=np.float64, + ) + + @cached_property + def uniform_sampling_mass(self) -> float: + """Return the summed systems-per-frame weights of the native streams.""" + if self.mixed_systems: + return float(self.nsystems / self.batch_size[0]) + return float(uniform_sampling_weights(self.mean_batch_nframes).sum()) + + @property + def uniform_nominal_mass(self) -> float: + """Return the systems-per-frame weight averaged over native passes.""" + return self.uniform_sampling_mass def get_batch(self, sys_idx: int | None = None) -> dict: # batch generation style altered by Ziyao Li: @@ -479,6 +520,63 @@ def get_batch(self, sys_idx: int | None = None) -> dict: b_data = self.get_batch_mixed() return b_data + def state_dict(self) -> dict[str, Any]: + """Return logical sampling progress for all systems. + + The caller must save the shared ``dp_random`` state separately, + once for all training and validation data sources. + """ + return { + "version": 1, + "batch_size": [int(size) for size in self.batch_size], + "sys_probs": np.asarray(self.sys_probs).tolist(), + "mixed_systems": self.mixed_systems, + "pick_idx": int(self.pick_idx), + "systems": [data.state_dict() for data in self.data_systems], + } + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore progress without materializing any frame data.""" + if state.get("version") != 1: + raise ValueError("Unsupported DeepmdDataSystem progress version") + if state.get("batch_size") != [int(size) for size in self.batch_size]: + raise ValueError("DeepmdDataSystem progress has different batch sizes") + try: + saved_probs = np.array(state.get("sys_probs"), dtype=float, copy=True) + except (TypeError, ValueError) as error: + raise ValueError( + "DeepmdDataSystem progress has invalid probabilities" + ) from error + if ( + saved_probs.shape != (self.nsystems,) + or not np.all(np.isfinite(saved_probs)) + or np.any(saved_probs < 0) + or not np.allclose(saved_probs, self.sys_probs, rtol=1e-12, atol=0) + ): + raise ValueError("DeepmdDataSystem progress has different probabilities") + if state.get("mixed_systems") != self.mixed_systems: + raise ValueError("DeepmdDataSystem progress has a different batching mode") + states = state.get("systems", []) + if len(states) != self.nsystems: + raise ValueError("DeepmdDataSystem progress has different systems") + pick_idx = state.get("pick_idx") + if type(pick_idx) is not int or not 0 <= pick_idx < self.nsystems: + raise ValueError("Invalid DeepmdDataSystem progress pick_idx") + # Validate every system before changing any reader's progress. + for data, data_state in zip(self.data_systems, states, strict=True): + data._validate_state_dict(data_state) + for data, data_state in zip(self.data_systems, states, strict=True): + data.load_state_dict(data_state) + # Tolerate recomputation roundoff, but retain the checkpoint's exact + # selection thresholds so the restored RNG resumes the same stream. + # Preserve an exactly matching dtype too: NumPy uses it to choose + # the tolerance for checking whether probabilities sum to one. + if np.array_equal(saved_probs, self.sys_probs): + saved_probs = np.array(self.sys_probs, copy=True) + self.sys_probs = saved_probs + self.pick_idx = pick_idx + self.__dict__.pop("test_data", None) + def get_batch_standard(self, sys_idx: int | None = None) -> dict: """Get a batch of data from the data systems in the standard way. @@ -497,7 +595,6 @@ def get_batch_standard(self, sys_idx: int | None = None) -> dict: if sys_idx is not None: self.pick_idx = sys_idx else: - # prob = self._get_sys_probs(sys_probs, auto_prob_style) self.pick_idx = dp_random.choice( np.arange(self.nsystems, dtype=np.int32), p=self.sys_probs ) @@ -585,6 +682,15 @@ def _merge_batch_data(self, batch_data: list[dict]) -> dict: ) b_data[kk] = merged.reshape(len(batch_data), -1) elif not vv["atomic"]: + if vv.get("special_shape") == "frame_major": + extents = {bb[kk].shape[1:] for bb in batch_data} + if len(extents) > 1: + raise ValueError( + f"The frame-major data {kk} has inconsistent " + f"extents {sorted(extents)} across systems; " + "mixed batching requires the same extent in " + "every system" + ) b_data[kk] = np.concatenate([bb[kk] for bb in batch_data], axis=0) else: b_data[kk] = np.zeros( @@ -668,6 +774,10 @@ def get_batch_pass_length(self) -> int: return compute_total_numb_batch(self.nbatches, self.sys_probs) + def get_epoch_batch_counts(self, world_size: int) -> np.ndarray: + """Return optimizer workloads for the replicated NumPy data reader.""" + return np.asarray(self.nbatches, dtype=np.float64) / world_size + def print_summary(self, name: str) -> None: print_summary( name, @@ -700,24 +810,6 @@ def _make_auto_ts(self, percent: float) -> list[int]: return ts - def _check_type_map_consistency( - self, type_map_list: list[list[str] | None] - ) -> list[str]: - ret = [] - for ii in type_map_list: - if ii is not None: - min_len = min([len(ii), len(ret)]) - for idx in range(min_len): - if ii[idx] != ret[idx]: - raise RuntimeError( - f"Inconsistent type map: {ret!s} {ii!s} in different data systems. " - "If you didn't set model/type_map, please set it, " - "since the type map of the model cannot be decided by data." - ) - if len(ii) > len(ret): - ret = ii - return ret - def _format_name_length(name: str, width: int) -> str: if len(name) <= width: @@ -820,115 +912,22 @@ def print_summary( ) -def process_sys_probs(sys_probs: list[float], nbatch: int) -> np.ndarray: - sys_probs = np.array(sys_probs) - type_filter = sys_probs >= 0 - assigned_sum_prob = np.sum(type_filter * sys_probs) - # 1e-8 is to handle floating point error; See #1917 - assert assigned_sum_prob <= 1.0 + 1e-8, ( - "the sum of assigned probability should be less than 1" - ) - rest_sum_prob = 1.0 - assigned_sum_prob - if not np.isclose(rest_sum_prob, 0): - rest_nbatch = (1 - type_filter) * nbatch - rest_prob = rest_sum_prob * rest_nbatch / np.sum(rest_nbatch) - ret_prob = rest_prob + type_filter * sys_probs - else: - ret_prob = sys_probs - assert np.isclose(np.sum(ret_prob), 1), "sum of probs should be 1" - return ret_prob - - -def prob_sys_size_ext(keywords: str, nsystems: int, nbatch: int) -> list[float]: - block_str = keywords.split(";")[1:] - block_stt = [] - block_end = [] - block_weights = [] - for ii in block_str: - stt = int(ii.split(":")[0]) - end = int(ii.split(":")[1]) - weight = float(ii.split(":")[2]) - assert weight >= 0, "the weight of a block should be no less than 0" - block_stt.append(stt) - block_end.append(end) - block_weights.append(weight) - nblocks = len(block_str) - block_probs = np.array(block_weights) / np.sum(block_weights) - sys_probs = np.zeros([nsystems], dtype=np.float64) - for ii in range(nblocks): - nbatch_block = nbatch[block_stt[ii] : block_end[ii]] - tmp_prob = [float(i) for i in nbatch_block] / np.sum(nbatch_block) - sys_probs[block_stt[ii] : block_end[ii]] = tmp_prob * block_probs[ii] - return sys_probs - - -def process_systems( - systems: str | list[str], patterns: list[str] | None = None -) -> list[str]: - """Process the user-input systems. - - If it is a single directory, search for all the systems in the directory. - If it is a list, each item in the list is treated as a directory to search. - If it is a single LMDB path, return it directly without expansion. - Check if the systems are valid. - - Parameters - ---------- - systems : str or list of str - The user-input systems - patterns : list of str, optional - The patterns to match the systems, by default None - - Returns - ------- - result_systems: list of str - The valid systems - """ - from deepmd.dpmodel.utils.lmdb_data import ( - is_lmdb, - ) - - # LMDB path: return directly without expansion - if isinstance(systems, str) and is_lmdb(systems): - return [systems] - - # Normalize input to a list of paths to search - if isinstance(systems, str): - search_paths = [systems] - elif isinstance(systems, list): - search_paths = systems - else: - # Handle unsupported input types - raise ValueError( - f"Invalid systems type: {type(systems)}. Must be str or list[str]." - ) - - # Iterate over the search_paths list and apply expansion logic to each path - result_systems = [] - for path in search_paths: - if patterns is None: - expanded_paths = expand_sys_str(path) - else: - expanded_paths = rglob_sys_str(path, patterns) - - result_systems.extend(expanded_paths) - - return result_systems - - def get_data( - jdata: dict[str, Any], + jdata: DataConfig, rcut: float, type_map: list[str] | None, modifier: Any | None, multi_task_mode: bool = False, -) -> DeepmdDataSystem: - """Get the data system. + *, + lmdb_factory: Callable[[str, int | str, str | None], Any] | None = None, + seed: int | None = None, +) -> Any: + """Build native data readers and compose independently configured groups. Parameters ---------- jdata - The json data + A data dictionary or a nonempty list of group dictionaries. rcut The cut-off radius, not used type_map @@ -937,33 +936,121 @@ def get_data( The data modifier multi_task_mode If in multi task mode + lmdb_factory + Optional backend adapter accepting an LMDB path, batch size, and + internal auto-probability rule. Without it, only NPY/HDF5 sources + are discovered. + seed + Seed for group selection. If omitted, draws from the shared training + random generator seed the group-selection streams. Returns ------- - DeepmdDataSystem - The data system + object + A native data system for a single source, or a grouped data system + with the same batch, statistics, and sampling-metadata interface. """ - systems = jdata["systems"] - rglob_patterns = jdata.get("rglob_patterns") - systems = process_systems(systems, patterns=rglob_patterns) - - batch_size = jdata["batch_size"] - sys_probs = jdata.get("sys_probs") - auto_prob = jdata.get("auto_prob", "prob_sys_size") - optional_type_map = not multi_task_mode - - data = DeepmdDataSystem( - systems=systems, - batch_size=batch_size, - test_size=1, # to satisfy the old api - shuffle_test=True, # to satisfy the old api - rcut=rcut, - type_map=type_map, - optional_type_map=optional_type_map, - modifier=modifier, - trn_all_set=True, # sample from all sets - sys_probs=sys_probs, - auto_prob_style=auto_prob, + from deepmd.utils.data_group import ( + GroupedDataSystem, ) - return data + groups = resolve_data_groups( + jdata, formats=("deepmd", "lmdb") if lmdb_factory is not None else ("deepmd",) + ) + needs_composition = len(groups) > 1 or any( + len(group.systems) > 1 and any(is_lmdb(path) for path in group.systems) + for group in groups + ) + seeds = ( + iter(data_child_seeds(seed, len(groups) + 1)) if needs_composition else iter(()) + ) + group_sources = [] + for group in groups: + sources = resolve_data_sources(group) + has_lmdb = any(source.format == "lmdb" for source in sources) + children = [] + for source in sources: + if source.format == "deepmd": + child = DeepmdDataSystem( + systems=list(source.systems), + batch_size=source.batch_size, + test_size=1, + shuffle_test=True, + rcut=rcut, + type_map=type_map, + optional_type_map=not multi_task_mode, + modifier=modifier, + trn_all_set=True, + sys_probs=group.sys_probs if not has_lmdb else None, + auto_prob_style=group.auto_prob or "prob_sys_size", + ) + else: + assert lmdb_factory is not None + child = lmdb_factory( + source.systems[0], source.batch_size, group.auto_prob + ) + children.append(child) + group_sources.append(children) + + sources = [source for children in group_sources for source in children] + numpy_sources = [ + source for source in sources if isinstance(source, DeepmdDataSystem) + ] + if len({source.mixed_type for source in numpy_sources}) > 1: + raise ValueError("All NumPy data systems must agree on mixed_type.") + ntypes = max(source.get_ntypes() for source in sources) + for source in numpy_sources: + source.sys_ntypes = ntypes + source.natoms_vec = [ + system.get_natoms_vec(ntypes).astype(int) for system in source.data_systems + ] + datasets = [ + ( + children[0] + if len(children) == 1 + else GroupedDataSystem( + children, + [None] * len(children), + seed=next(seeds), + uniform=group.auto_prob == "prob_uniform", + ) + ) + for group, children in zip(groups, group_sources, strict=True) + ] + return ( + datasets[0] + if len(datasets) == 1 + else GroupedDataSystem( + datasets, [group.prob for group in groups], seed=next(seeds) + ) + ) + + +def get_neighbor_stat_data(jdata: DataConfig, type_map: list[str] | None) -> Any: + """Expose every discovered source to backend neighbor-statistic kernels.""" + sources = [] + for group in resolve_data_groups(jdata): + for source in resolve_data_sources(group): + if source.format == "deepmd": + sources.append( + DeepmdDataSystem( + list(source.systems), + source.batch_size, + 1, + type_map=type_map, + trn_all_set=True, + ) + ) + else: + from deepmd.dpmodel.utils.lmdb_data import ( + make_neighbor_stat_data, + ) + + sources.append(make_neighbor_stat_data(source.systems[0], type_map)) + return SimpleNamespace( + data_systems=[system for source in sources for system in source.data_systems], + system_dirs=[path for source in sources for path in source.system_dirs], + mixed_type=any(source.mixed_type for source in sources), + get_ntypes=lambda: max(source.get_ntypes() for source in sources), + get_batch=lambda: None, + ) diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py index 5dd3623cba..2cf7d82fd4 100644 --- a/deepmd/utils/finetune.py +++ b/deepmd/utils/finetune.py @@ -221,17 +221,25 @@ def _native_spin_types( def _validate_native_spin_finetune( target_config: Mapping[str, Any], pretrained_config: Mapping[str, Any], -) -> None: - """Validate the native-spin transfer contract before model construction.""" +) -> bool: + """Validate spin transfer and report whether a spin-free source needs conversion.""" target_types = _native_spin_types(target_config) if target_types is None: - return + return False pretrained_types = _native_spin_types(pretrained_config) if pretrained_types is None: + if pretrained_config.get("spin") is None: + if target_config["type_map"] != pretrained_config["type_map"]: + raise ValueError( + "Spin-free to native-spin fine-tuning requires identical ordered " + f"type maps: pretrained={pretrained_config['type_map']}, " + f"target={target_config['type_map']}." + ) + return True raise ValueError( - "Native-spin fine-tuning requires a native-spin pretrained model; " - "automatic conversion from a spin-free model is not supported." + "Native-spin fine-tuning requires a spin-free or native-spin " + "pretrained model." ) if pretrained_types and pretrained_types != target_types: raise ValueError( @@ -239,6 +247,86 @@ def _validate_native_spin_finetune( f"fine-tuning is not supported: pretrained={sorted(pretrained_types)}, " f"target={sorted(target_types)}." ) + return False + + +def _learned_model_config( + model: Mapping[str, Any], +) -> Mapping[str, Any] | None: + """Find the unique descriptor-bearing part of a plain or composed model.""" + candidates = model["models"] if model.get("type") == "linear_ener" else (model,) + learned = [candidate for candidate in candidates if "descriptor" in candidate] + return learned[0] if len(learned) == 1 else None + + +def _inherit_native_spin_fitting_widths( + target: dict[str, Any], pretrained: Mapping[str, Any] +) -> None: + """Resolve automatic target widths against the pretrained fitting layout.""" + from deepmd.utils.argcheck import ( + descrpt_args_plugin, + fitting_args_plugin, + ) + + target_learned = _learned_model_config(target) + source_learned = _learned_model_config(pretrained) + if target_learned is None or source_learned is None: + return + descriptor_schema = descrpt_args_plugin.get_argument("dpa4c") + fitting_schema = fitting_args_plugin.get_argument("dpa4_ener") + descriptor_types = {descriptor_schema.name, *descriptor_schema.alias} + glu_types = {fitting_schema.name, *fitting_schema.alias} + fitting = target_learned.get("fitting_net", {}) + source_fitting = source_learned.get("fitting_net", {}) + if ( + target_learned["descriptor"].get("type") not in descriptor_types + or source_learned["descriptor"].get("type") not in descriptor_types + or fitting.get("type") not in glu_types + or source_fitting.get("type") not in glu_types + ): + return + neuron = fitting_schema.normalize_value(fitting).get("neuron") + target_widths = [0] if neuron is None else list(neuron) + if 0 not in target_widths: + return + source_fitting = fitting_schema.normalize_value(source_fitting) + source_neuron = source_fitting.get("neuron") + source_widths = [0] if source_neuron is None else list(source_neuron) + if len(target_widths) != len(source_widths): + raise ValueError( + "Automatic native-spin fitting widths require matching layer counts: " + f"pretrained={len(source_widths)}, target={len(target_widths)}." + ) + if 0 in source_widths: + from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, + ) + from deepmd.dpmodel.fitting.dpa4_ener import ( + _resolve_auto_neuron, + ) + + descriptor_params = descriptor_schema.normalize_value( + source_learned["descriptor"] + ) + descriptor_params.pop("type") + type_map = source_learned.get("type_map", pretrained["type_map"]) + descriptor_params.update(ntypes=len(type_map), type_map=list(type_map)) + descriptor = DescrptDPA4C(**descriptor_params) + source_widths = _resolve_auto_neuron( + source_widths, + dim_descrpt=descriptor.get_dim_out(), + numb_fparam=source_fitting.get("numb_fparam", 0), + numb_aparam=source_fitting.get("numb_aparam", 0), + dim_case_embd=source_fitting.get("dim_case_embd", 0), + case_film_embd=source_fitting.get("case_film_embd", False), + use_aparam_as_mask=source_fitting.get("use_aparam_as_mask", False), + ) + for alias in fitting_schema.sub_fields["neuron"].alias: + fitting.pop(alias, None) + fitting["neuron"] = [ + source_width if target_width == 0 else target_width + for source_width, target_width in zip(source_widths, target_widths, strict=True) + ] class FinetuneRuleItem: @@ -249,6 +337,7 @@ def __init__( model_branch: str = "Default", random_fitting: bool = False, resuming: bool = False, + convert_to_native_spin: bool = False, ) -> None: """ The rules for fine-tuning the model from pretrained model. @@ -265,12 +354,15 @@ def __init__( If true, the fitting net will be randomly initialized instead of inherit from the pretrained model. resuming If true, the model will just resume from model_branch without fine-tuning. + convert_to_native_spin + If true, initialize the native-spin model from a spin-free branch. """ self.p_type_map = p_type_map self.type_map = type_map self.model_branch = model_branch self.random_fitting = random_fitting self.resuming = resuming + self.convert_to_native_spin = convert_to_native_spin self.update_type = self.p_type_map != self.type_map def get_index_mapping(self) -> list[int]: @@ -398,7 +490,9 @@ def build(self) -> tuple[dict[str, Any], dict[str, FinetuneRuleItem]]: model_branch_from=model_branch_from, ) ) - finetune_rule.resuming = resuming + finetune_rule.resuming = ( + resuming and not finetune_rule.convert_to_native_spin + ) finetune_links[model_key] = finetune_rule return model_config, finetune_links @@ -444,7 +538,9 @@ def build_single_rule( model_branch_chosen = model_alias_dict[model_branch_chosen] single_config_chosen = deepcopy(model_dict_params[model_branch_chosen]) - _validate_native_spin_finetune(single_config, single_config_chosen) + convert_to_native_spin = _validate_native_spin_finetune( + single_config, single_config_chosen + ) old_type_map = single_config_chosen["type_map"] new_type_map = single_config["type_map"] @@ -453,6 +549,7 @@ def build_single_rule( type_map=new_type_map, model_branch=model_branch_chosen, random_fitting=new_fitting, + convert_to_native_spin=convert_to_native_spin, ) if self.change_model_params: self._apply_pretrained_model_params( @@ -462,6 +559,8 @@ def build_single_rule( model_branch_chosen=model_branch_chosen, random_fitting=new_fitting, ) + if convert_to_native_spin and not new_fitting: + _inherit_native_spin_fitting_widths(single_config, single_config_chosen) return single_config, finetune_rule def _apply_pretrained_model_params( @@ -490,6 +589,18 @@ def _apply_pretrained_model_params( single_config["descriptor"] = deepcopy(pretrained_config["descriptor"]) if not random_fitting: single_config["fitting_net"] = deepcopy(pretrained_config["fitting_net"]) + # A long-range module reads the descriptor beside the fitting and + # is carried over with it; the trainability set in the input + # applies to its parts as to the descriptor and the fitting. + if "long_range" in pretrained_config: + long_range = deepcopy(pretrained_config["long_range"]) + for part in ("charge_net", "kernel"): + long_range.setdefault(part, {})["trainable"] = ( + single_config.get("long_range", {}) + .get(part, {}) + .get("trainable", True) + ) + single_config["long_range"] = long_range log.info( "Change the '%s' model configurations according to pretrained branch '%s'.", model_branch, diff --git a/deepmd/utils/json_schema.py b/deepmd/utils/json_schema.py index 1a080ec777..6430bbff8a 100644 --- a/deepmd/utils/json_schema.py +++ b/deepmd/utils/json_schema.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""JSON Schema for model inputs before preset expansion.""" +"""JSON Schema for model presets and grouped data input forms.""" import json import re @@ -14,6 +14,24 @@ _REGIONS = ("type_map", "descriptor", "fitting_net") +def with_data_groups(schema: dict[str, Any], multi_task: bool) -> dict[str, Any]: + """Describe the same training-data fields for dictionaries and list items.""" + training = schema["properties"]["training"] + if multi_task: + training = training["properties"]["data_dict"]["additionalProperties"] + data = training["properties"]["training_data"] + data["items"] = { + "type": "object", + "properties": data["properties"], + "required": data["required"], + } + if "allOf" in data: + data["items"]["allOf"] = data["allOf"] + data["minItems"] = 1 + data["properties"]["prob"].update(minimum=0, maximum=1) + return schema + + def _preset_pattern(names: list[str]) -> str: """Match complete preset names case-insensitively with ECMA-262 syntax.""" alternatives = [ diff --git a/deepmd/utils/model_preset_data.py b/deepmd/utils/model_preset_data.py index 18df171591..e1274bb7f6 100644 --- a/deepmd/utils/model_preset_data.py +++ b/deepmd/utils/model_preset_data.py @@ -124,7 +124,7 @@ "descriptor": { "channels": 128, "lmax": 6, - "n_blocks": 10, + "n_blocks": 8, "mixing_layers": 4, "radial_so2_rank": 4, "n_focus": 3, diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py index 109de19fd5..ac927d0e4f 100644 --- a/deepmd/utils/model_stat.py +++ b/deepmd/utils/model_stat.py @@ -99,7 +99,8 @@ def _collect_stat_batches( if "natoms_vec" in stat_data: stat_data["natoms_vec"] = stat_data["natoms_vec"].astype(np.int32) batch = normalize_batch(stat_data) - frame_mask = pair_margin_frame_mask(batch, min_pair_dist > 0.0) + filter_enabled = getattr(stat_data, "pair_filter_enabled", min_pair_dist > 0.0) + frame_mask = pair_margin_frame_mask(batch, filter_enabled) if frame_mask is not None and not frame_mask.all(): if not frame_mask.any(): continue diff --git a/deepmd/utils/path.py b/deepmd/utils/path.py index 44bbe34edf..c511ea52fb 100644 --- a/deepmd/utils/path.py +++ b/deepmd/utils/path.py @@ -22,6 +22,7 @@ Self, ) from wcmatch.glob import ( + GLOBSTAR, globfilter, ) @@ -393,7 +394,7 @@ def glob(self, pattern: str) -> list["DPPath"]: ] return [ type(self)(f"{self.root_path}#{pp}", mode=self.mode) - for pp in globfilter(subpaths, self._connect_path(pattern)) + for pp in globfilter(subpaths, self._connect_path(pattern), flags=GLOBSTAR) ] def rglob(self, pattern: str) -> list["DPPath"]: @@ -410,7 +411,7 @@ def rglob(self, pattern: str) -> list["DPPath"]: list[DPPath] list of paths """ - return self.glob("**" + pattern) + return self.glob("**/" + pattern) @property def _keys(self) -> list[str]: diff --git a/deepmd/utils/preset_out_bias.py b/deepmd/utils/preset_out_bias.py index 48a53f644a..f5bc55aaed 100644 --- a/deepmd/utils/preset_out_bias.py +++ b/deepmd/utils/preset_out_bias.py @@ -401,6 +401,48 @@ def preset_out_bias_rows( return rows +def recover_nonfinite_preset_bias( + preset_out_bias: dict[str, list[list | None]], + type_map: list[str], + stored_bias: np.ndarray, + keys: list[str], + sizes: list[int], +) -> dict[str, np.ndarray]: + """Recover non-finite stored cells from presets without fixing fitted outputs. + + Parameters + ---------- + preset_out_bias + Normalized checkpoint preset. Unassigned cells cannot be recovered. + type_map + Element names in stored type order, including types absent from data. + stored_bias + Stored bias with shape (n_out, ntypes, max_size). + keys + Output names in stored order. + sizes + Flattened size of each output. + + Returns + ------- + dict[str, np.ndarray] + Changed output rows. All finite stored cells are preserved. + """ + preset_rows = preset_out_bias_rows( + preset_out_bias, type_map, [], stored_bias, keys, sizes, keep_unassigned=True + ) + recovered = {} + for index, (key, size) in enumerate(zip(keys, sizes, strict=True)): + if key not in preset_rows: + continue + stored = stored_bias[index, :, :size] + replacement = preset_rows[key] + mask = ~np.isfinite(stored) & np.isfinite(replacement) + if np.any(mask): + recovered[key] = np.where(mask, replacement, stored) + return recovered + + def override_assigned_bias( bias: np.ndarray, assigned_bias: np.ndarray | None, diff --git a/deepmd/utils/preset_out_bias_tables.json b/deepmd/utils/preset_out_bias_tables.json index b21a05811e..1457deb819 100644 --- a/deepmd/utils/preset_out_bias_tables.json +++ b/deepmd/utils/preset_out_bias_tables.json @@ -472,5 +472,107 @@ "Pb": -5245.360635715531, "Bi": -5838.120614817622 } + }, + "openmd": { + "source": "OpenMD ABACUS PBE PW atomic reference energies in eV; ABACUS_PBE_PW_atom_E0.json, H-Bi (83 elements), ABACUS 3.9.0 !FINAL_ETOT_IS; non-spin-polarized, 100 Ry cutoff, 30 Angstrom cell, Gaussian smearing 0.01 Ry; La uses an accepted unconverged final-step energy", + "values": { + "H": -12.560490553805243, + "He": -75.93894279525861, + "Li": -197.02706138891466, + "Be": -380.9398767190545, + "B": -73.46336926845866, + "C": -154.90244724923775, + "N": -262.67292484125375, + "O": -443.0343602483389, + "F": -666.9599186786108, + "Ne": -834.8469742136249, + "Na": -1242.7178587090743, + "Mg": -25.898184583845225, + "Al": -59.456590941108324, + "Si": -109.93805025060193, + "P": -3054.572178879562, + "S": -296.4442167325705, + "Cl": -444.565461150324, + "Ar": -577.5481371256309, + "K": -816.8902908164619, + "Ca": -1051.506833490634, + "Sc": -1312.493731843498, + "Ti": -1608.7058057558409, + "V": -2000.4614481107299, + "Cr": -2271.0072192249995, + "Mn": -2955.314766598183, + "Fe": -3212.705379389896, + "Co": -3680.699255922714, + "Ni": -4560.594847227876, + "Cu": -4965.331239170799, + "Zn": -1730.1064697876304, + "Ga": -2078.1744226015844, + "Ge": -2426.261793652987, + "As": -2881.7920489669164, + "Se": -3295.120612744489, + "Br": -426.6432826063096, + "Kr": -505.82759543327364, + "Rb": -735.9387663298128, + "Sr": -929.9320319807304, + "Y": -1105.7623041248487, + "Zr": -1370.4398149568858, + "Nb": -1685.5929111214323, + "Mo": -1901.9732811740478, + "Tc": -2184.1104857459873, + "Ru": -2742.0353598251418, + "Rh": -3185.4184768339765, + "Pd": -3683.105491676735, + "Ag": -3906.989494408327, + "Cd": -4228.894062306103, + "In": -1805.8559760205173, + "Sn": -2140.8564151271016, + "Sb": -2471.5330180822257, + "Te": -285.31701318449774, + "I": -350.7009208910242, + "Xe": -510.3075936830297, + "Cs": -635.6785167026409, + "Ba": -759.9931396990479, + "La": -950.8377156904378, + "Ce": -967.8038269507553, + "Pr": -998.0022169607034, + "Nd": -1035.4427966959247, + "Pm": -1090.7678703417816, + "Sm": -1085.726379159036, + "Eu": -1114.254988887816, + "Gd": -1127.6701518356654, + "Tb": -1156.031115926042, + "Dy": -1147.8580058338957, + "Ho": -1197.006249013279, + "Er": -1224.0375204753566, + "Tm": -1182.6392327873907, + "Yb": -2527.830303042947, + "Lu": -1282.950916323345, + "Hf": -1528.0335315810078, + "Ta": -1762.1544789342358, + "W": -1854.385401723439, + "Re": -2350.6242939579165, + "Os": -2742.9296928209546, + "Ir": -3125.1050363155114, + "Pt": -3547.7848701247403, + "Au": -3734.78716520503, + "Hg": -4221.467752095196, + "Tl": -1303.0098606042245, + "Pb": -1911.8667279529952, + "Bi": -6345.199411321936 + } + }, + "domain_cluster": { + "source": "Domain_cluster atomic energy reference biases in eV; user-provided nine-element table", + "values": { + "Al": -52.96613964130058, + "Si": -102.54342787198362, + "Ni": -4601.828023238278, + "Cu": -1306.1758472225818, + "Ru": -2570.6691516730557, + "Pd": -3458.0089798503986, + "Ag": -1005.129093525484, + "Pt": -3264.8411944198538, + "Au": -901.9377033597224 + } } } diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index 8710693607..55e23b22f2 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -76,4 +76,18 @@ def shuffle(x: np.ndarray) -> None: _RANDOM_GENERATOR.shuffle(x) -__all__ = ["choice", "random", "seed", "shuffle"] +def get_state() -> tuple: + """Return the generator state using checkpoint-safe Python values.""" + kind, keys, pos, has_gauss, cached_gaussian = _RANDOM_GENERATOR.get_state() + return (kind, keys.tolist(), int(pos), int(has_gauss), float(cached_gaussian)) + + +def set_state(state: tuple) -> None: + """Restore a state returned by :func:`get_state`.""" + kind, keys, pos, has_gauss, cached_gaussian = state + _RANDOM_GENERATOR.set_state( + (kind, np.asarray(keys, dtype=np.uint32), pos, has_gauss, cached_gaussian) + ) + + +__all__ = ["choice", "get_state", "random", "seed", "set_state", "shuffle"] diff --git a/doc/backend.md b/doc/backend.md index c5155cc991..38b687eab8 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -77,8 +77,7 @@ 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 +`edge_vec` ABI (`lower_input_kind: edge_vec`). A graph model frozen with `dp --pt-expt freeze --lower-kind graph` uses the NeighborGraph ABI (`lower_input_kind: graph`). The `dp --pt-expt compress` workflow can instead export a compressed DPA-1 model through the compact canonical diff --git a/doc/env.md b/doc/env.md index ba08490a06..b0b5280757 100644 --- a/doc/env.md +++ b/doc/env.md @@ -63,11 +63,22 @@ Inference batch size, calculated by multiplying the number of frames with the nu Setting this variable disables automatic growth; out-of-memory errors can still reduce the batch size. -With the native PyTorch CUDA allocator, automatic multi-frame evaluation starts -with one calibration frame and limits growth using the observed workspace and -available device memory, leaving a 10% margin. This policy applies to both -`dp test` and full validation in the PyTorch and PyTorch Exportable backends. -Other GPU allocators and backends grow batches until an out-of-memory error. +On GPUs, automatic sizing grows successful full batches exponentially and reuses +that atom budget across systems. An out-of-memory error halves the actual +attempted batch size and retries the same input. Growth beyond the failed size +is retried after two successful full batches; repeated recovery probes +increase the interval to four, then eight batches. This short, bounded delay +favors recovering larger batches when subsequent input needs less memory, at +the cost of occasional OOM retries even on a steady workload. Short final +batches neither shrink the learned budget nor trigger growth. + +This policy applies to both `dp test` and full validation, without profiling CUDA +allocator state on successful inference. It does not search for an exact maximum +or guarantee the throughput-optimal batch size. Explicitly setting a positive +`DP_INFER_BATCH_SIZE` disables growth, including these recovery probes. +Hook-based feature extraction that retries the whole input defers recovery +probes until a complete evaluation succeeds, so a known OOM cannot cause an +endless restart of that input. ::: :::{envvar} DP_BACKEND diff --git a/doc/inference/cxx.md b/doc/inference/cxx.md index 4af32776ff..b31a9274a3 100644 --- a/doc/inference/cxx.md +++ b/doc/inference/cxx.md @@ -62,6 +62,7 @@ The backend plugin search order is: If the requested plugin or its backend runtime cannot be loaded, only that backend fails with a `Unable to load ... backend plugin` error. Other backends can still run as long as their own plugins and runtime libraries are available. +Install the plugins from the same build as `libdeepmd_cc`. The entry points a plugin exports carry the version of the backend interface it was built against, so `libdeepmd_cc` finds none of them in a plugin of another interface version, and that backend fails to load in the same way. This also allows a no-backend `libdeepmd_cc` or `libdeepmd_c` build; install or copy the backend plugin next to the library, or set {envvar}`DP_BACKEND_PLUGIN_PATH`, before loading a model that uses that backend. Build instructions for this layout are in [Install DeePMD-kit's C++ interface](../install/install-from-source.md#install-deepmd-kits-c-interface). diff --git a/doc/install/install-from-source.md b/doc/install/install-from-source.md index f07bfd2060..67a998dccd 100644 --- a/doc/install/install-from-source.md +++ b/doc/install/install-from-source.md @@ -180,7 +180,7 @@ Build CPU variant or GPU variant with CUDA or ROCM support. **Type**: Path; **Default**: Detected automatically -The path to the CUDA toolkit directory. CUDA 9.0 or later is supported. NVCC is required. +The path to the CUDA toolkit directory. CUDA 9.0 or later is supported, and CUDA 11.8 or later for the PyTorch backend. NVCC is required. ::: :::{envvar} ROCM_ROOT @@ -518,7 +518,7 @@ If `TRUE`, Build GPU support with CUDA toolkit. **Type**: `PATH`, **Default**: [Search automatically](https://cmake.org/cmake/help/latest/module/FindCUDAToolkit.html) -The path to the CUDA toolkit directory. CUDA 9.0 or later is supported. NVCC is required. +The path to the CUDA toolkit directory. CUDA 9.0 or later is supported, and CUDA 11.8 or later for the PyTorch backend. NVCC is required. See also [CMake documentation](https://cmake.org/cmake/help/latest/module/FindCUDAToolkit.html). ::: diff --git a/doc/install/install-lammps.md b/doc/install/install-lammps.md index 30828e614e..8b63b56b4f 100644 --- a/doc/install/install-lammps.md +++ b/doc/install/install-lammps.md @@ -1,8 +1,8 @@ # Install LAMMPS An AI agent can build LAMMPS with the DeePMD-kit module and Kokkos. DPA4/SeZM -uses `pair_style deepmd/kk`; DPA4C uses `pair_style dpa4spin/kk`. See -[Install with an AI agent](install-with-agent.md). +and DPA4C use `pair_style deepmd/kk`; compressed native-spin DPA4C models use +`pair_style dpa4spin/kk`. See [Install with an AI agent](install-with-agent.md). There are two ways to install LAMMPS: the built-in mode and the plugin mode. The built-in mode builds LAMMPS along with the DeePMD-kit and DeePMD-kit will be loaded automatically when running LAMMPS. The plugin mode builds LAMMPS and a plugin separately, so one needs to use `plugin load` command to load the DeePMD-kit's LAMMPS plugin library. @@ -84,8 +84,8 @@ make -j4 make install ``` -To build the GPU-resident `pair_style deepmd/kk`, enable the LAMMPS Kokkos -package. The pair style uses the DeePMD C API: +To build the GPU-resident pair styles `deepmd/kk` and `dpa4spin/kk`, enable the +LAMMPS Kokkos package. The pair styles use the DeePMD C API: ```bash cmake -D PKG_KOKKOS=ON \ @@ -129,9 +129,9 @@ make install Configure LAMMPS with `PKG_KOKKOS=ON`, then configure the DeePMD plugin with `DEEPMD_LAMMPS_KOKKOS=ON` and `Kokkos_DIR` pointing to the Kokkos package -exported by the LAMMPS build. This registers the GPU-resident -`pair_style deepmd/kk`; the plugin and built-in integrations use the same C API -device-graph interface. +exported by the LAMMPS build. This registers the GPU-resident pair styles +`deepmd/kk` and `dpa4spin/kk`; the plugin and built-in integrations use the same +C API device-graph interface. If everything works fine, you will end up with an executable `${deepmd_root}/bin/lmp`. @@ -145,3 +145,23 @@ ${deepmd_root}/bin/lmp -h > ```sh > patchelf --add-rpath "${tensorflow_root}/lib" liblammps.so > ``` + +## MPI for runs on several GPUs + +With the KOKKOS package, each MPI rank of LAMMPS uses one GPU and keeps its atom data in GPU memory, and the ranks exchange atom data at every step. Build LAMMPS and the DeePMD-kit C++ interface against the same MPI library, one that can access GPU memory (CUDA-aware MPI), such as Open MPI with UCX. When LAMMPS detects an MPI library without CUDA support, such as the MPICH package of conda-forge, it turns GPU-aware communication off and copies the positions and forces through host memory at every step. An MPI library that LAMMPS cannot identify is assumed to support CUDA; if it does not, run LAMMPS with `-pk kokkos gpu/aware off`. + +The Open MPI package of conda-forge is built with CUDA support but disables it by default; `ompi_info --all | grep opal_built_with_cuda_support` shows whether an Open MPI installation was built with it. Install the package in an environment of its own: + +```bash +conda create -n openmpi -c conda-forge openmpi ucx +``` + +Configure LAMMPS and the DeePMD-kit C++ interface in new build directories with both compiler wrappers of that environment, `-D MPI_C_COMPILER=${mpi_root}/bin/mpicc -D MPI_CXX_COMPILER=${mpi_root}/bin/mpicxx`, where `${mpi_root}` is its prefix; CMake finds the MPI of each language separately. Launch a run on several GPUs with the `mpirun` of that installation, enabling its CUDA support for the launch: + +```bash +mpirun --mca opal_cuda_support 1 -x UCX_MEMTYPE_CACHE=n -np 2 lmp -k on g 2 -sf kk -in in.lammps +``` + +With `UCX_MEMTYPE_CACHE=n`, UCX queries the memory type of every buffer instead of taking it from a cache that can report a GPU buffer as host memory and make UCX abort. + +Launch a run on a single MPI rank without these options: it exchanges no atom data, and the CUDA support adds a memory-type query to every MPI call, which can lengthen the step of a small system by more than 10%. diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index ad70c54b6f..b3d9e0033b 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -78,10 +78,9 @@ unnecessary and not recommended (see [Hardware selection](#hardware-selection)). > descriptors. On the conservative **energy** path it is only an initial > neighbor-search capacity that grows on demand, so it never truncates the > neighbor list and you do not need to size it to the true maximum neighbor count. -> The native spin scheme shares this energy path, so it grows on demand too; only -> the denoising (`dens`) path and the `deepspin` spin scheme cap the list at -> `sum(sel)`. You can also set `sel` to `auto` or `auto:factor` to size it from -> the training data. +> The native spin scheme shares this energy path, so it grows on demand too; +> the denoising (`dens`) path caps the list at `sum(sel)`. You can also set `sel` +> to `auto` or `auto:factor` to size it from the training data. ### Presets @@ -231,23 +230,14 @@ default training path. See `examples/water/dpa4/input_dens.json` for an example. ### Spin DPA4/SeZM supports the DeePMD-kit spin convention through the standard -`model.spin` block. Two schemes are available, selected by `model.spin.scheme`: - -- `native` — the per-atom spin vector enters the descriptor as an - equivariant feature, and the magnetic force is the spin gradient of the - energy. No virtual atoms are introduced, so the neighbor list and type map - keep their real-system sizes. -- `deepspin` (default) — the classical DeepSpin representation, in which each - magnetic atom is paired with a virtual atom displaced along its spin. It is - the default, so a `model.spin` block without an explicit `scheme` reproduces - the classical behaviour shared with every non-SeZM spin model. - -Both schemes train against the conservative `ener_spin` loss, share the same +`model.spin` block. The `native` scheme passes the per-atom spin vector into the +descriptor as an equivariant feature, and the magnetic force is the negative +spin gradient of the energy. The neighbor list and type map keep their +real-system sizes. + +Spin models train against the conservative `ener_spin` loss, use the `spin` / `force_mag` data convention, and are not combined with the `dens` mode. -Because the dataset and loss are identical, switching `scheme` does not require -any change to the data or the loss block. Complete inputs are in -`examples/spin/dpa4/`: `input.json` for the native scheme and -`input-deepspin.json` for the deepspin scheme. See +A complete input is provided in `examples/spin/dpa4/input.json`. See [training spin energy models](train-energy-spin.md) for the general workflow. `use_spin` accepts a per-type boolean list or, to avoid enumerating a large @@ -294,47 +284,11 @@ the magnetic force come from a single energy gradient: where $\mathbf{s}_i$ is the input spin vector. The model does not rescale the spin internally, so $\mathbf{s}_i$ keeps the dataset's `spin` convention and the -magnetic force is reported in the matching units of `force_mag` -- there is no -`virtual_scale` factor in this scheme. The magnetic force is reported on the -magnetic atom types only, matching the `force_mag` label. The native scheme +magnetic force is reported in the matching units of `force_mag`, on magnetic +atom types only, matching the `force_mag` label. The native scheme relies on the descriptor's angular degrees to represent the spin direction, so it requires `lmax >= 1` (the default). -#### DeepSpin virtual-atom scheme (default) - -```json -{ - "model": { - "type": "dpa4", - "type_map": [ - "Ni", - "O" - ], - "spin": { - "use_spin": [ - true, - false - ], - "virtual_scale": [ - 0.314 - ], - "scheme": "deepspin" - }, - "descriptor": { - "sel": 120, - "rcut": 6.0 - } - } -} -``` - -The `deepspin` scheme augments each magnetic atom with a virtual atom at -`coord + spin * virtual_scale`, which doubles the internal neighbor capacity and -type map. `virtual_scale` is required by this scheme and ignored by the native -scheme. The `deepspin` scheme cannot be combined with `bridging_method`: a -magnetic atom lies inside the bridging window of its own virtual partner, so a -bridged spin model uses the native scheme. - ### Multi-task / shared fitting DPA4/SeZM supports shared-fitting multitask training. With @@ -568,6 +522,15 @@ shipping built in. On other GPUs the kernels fall back to conservative configurations, and `dp --pt freeze` tunes the missing entries on the local GPU before exporting, a one-off sweep of a few minutes baked into the `.pt2`. +The frozen first-order FP32 SO(2) stack automatically uses edge cohorts for +the validated two-focus layout (`lmax=3`, 32 channels per focus, two gated +layers and focus competition). It completes the backward layers for up to +8192 edges before reusing scratch for the next cohort; saved activations and +returned gradients still cover all edges. Training, higher-order derivatives +and other layouts keep their existing paths. No separate feature switch is +required. This runtime launcher serves online compilation; the official +CPU-traced exports retain the native deployable backward. + `DP_CUTILE_INFER` replaces the whole SeZM edge pipeline — Wigner monomials, rotate-and-mix, the gated SO(2) mixing stack, the attention aggregation and the force / virial assembly — with kernels written in the `cuda.tile` DSL. On an @@ -741,11 +704,6 @@ Two different export routes produce one: The two are distinct ABIs consumed by different C++ paths; neither is a dense route, and `--lower-kind` selects between the pt_expt lowers only. -Which models lose multi-rank, on either route: - -- **The deepspin (virtual-atom) spin scheme**, which overrides the export ABI - to `nlist` because it expands virtual atoms inside the graph. - **Zone bridging is multi-rank capable** (including combined with native spin). The Source Freeze Propagation gate is a product over each node's full *outgoing*-edge set, which no single rank observes for ghost owners; the @@ -799,8 +757,7 @@ configured with `spin.scheme: native` (see [Native spin (magnetic)](#native-spin-magnetic) below) -- can be frozen through a NeighborGraph-native inference path instead of the legacy dense neighbor-list path. Frame-level charge/spin conditioning -(`add_chg_spin_ebd`) and zone bridging are graph-eligible too; only the -`deepspin` virtual-atom spin scheme remains dense-only: +(`add_chg_spin_ebd`) and zone bridging are graph-eligible too: ```bash dp --pt_expt freeze -o model.pt2 --lower-kind graph @@ -814,17 +771,12 @@ it considers every neighbor within `rcut` rather than the capacity fixed by `sel`. For the default conservative-energy path, `sel` is already only an initial search capacity that grows on demand (see the note in [Minimal input](#minimal-input)), so graph and dense results agree there down to -floating-point noise. Where `sel` does act as a hard cap -- the `dens` -denoising path, or the `deepspin` virtual-atom spin scheme, both of which -size the neighbor list to `sum(sel)` -- the graph route's larger neighbor set -can diverge from the dense route, the same deliberate divergence documented -for DPA-1/DPA-2. - -A DPA4/SeZM descriptor configured with `deepspin`-scheme spin is not -graph-eligible and always runs the dense route regardless of `--lower-kind`; -`--lower-kind graph` on such a model raises an error at freeze time instead -of exporting a silently-dense-only artifact. `native`-scheme spin and zone -bridging are the opposite case: they have *no* dense route at all -- +floating-point noise. On the `dens` denoising path, `sel` acts as a hard cap +that sizes the neighbor list to `sum(sel)`. The graph route's larger neighbor +set can therefore diverge from the dense route, the same deliberate divergence +documented for DPA-1/DPA-2. + +Native spin and zone bridging have no dense route at all -- the analytical bridging term has no dense injection site -- so they are always graph-frozen, `--lower-kind` notwithstanding. Charge/spin conditioning rides the graph lower and constrains neither. Note that a @@ -855,15 +807,8 @@ artifact and support multi-rank LAMMPS -- see [Native spin the `native`-scheme JSON example) is graph-only: a native-spin DPA4/SeZM model has no dense (nlist) lower at all, so it is always frozen through the NeighborGraph path regardless of `--lower-kind`, and always runs through it -at inference time. The `deepspin` virtual-atom scheme is unaffected and -keeps using the dense route described in [Spin](#spin) and [Multi-GPU (MPI) -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 - `deepspin`-scheme model raises an error at freeze time, per the dense/graph - eligibility rule above. +at inference time. + - **Graph route only, no dense fallback.** A native-spin descriptor has only the graph lower. This is not a spin-specific restriction on the CLI: every graph-capable DPA4/SeZM model, plain-energy included, is frozen to the @@ -1020,10 +965,8 @@ closed over the one-hop neighbor shell. completed across ranks by one reverse-accumulate plus one forward-broadcast border exchange. See [Multi-GPU (MPI) inference](#multi-gpu-mpi-inference). -- The pt_expt graph-native inference route is unavailable only for - `deepspin`-scheme spin, which stays on the dense route. Charge/spin - conditioning, zone bridging and `native`-scheme spin are all - graph-eligible. +- Charge/spin conditioning, zone bridging and native spin are all eligible + for the pt_expt graph-native inference route. - `spin.scheme: native` is graph-only (it has no dense route) and supports multi-rank LAMMPS: ghost node features ride `border_op` per interaction block and ghost spins arrive through the LAMMPS `sp` forward-comm. diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md index 7e79d8098e..2278d7e73d 100644 --- a/doc/model/dpa4c.md +++ b/doc/model/dpa4c.md @@ -332,24 +332,60 @@ now summed with. ## Model compression Compression is the deployment step. It replaces the analytic radial functions -and their type-pair modulation with tabulated splines evaluated by fused CUDA -kernels, and re-exports the model in the compact canonical graph form that the -fast inference path consumes. Because the radial map is analytically bounded and +and their type-pair modulation with tabulated splines evaluated by fused CPU or +CUDA kernels. Eligible CUDA models use the compact canonical graph form; +CPU models use the graph form with their dedicated kernels. Because the radial map is analytically bounded and vanishes at `rcut`, the table needs no extrapolation region and no overflow checking. -Train, freeze, then compress the frozen archive: +Freeze and compress a checkpoint in one operation: ```bash dp --pt-expt train input.json -dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph -dp --pt-expt compress -i frozen_model.pt2 -o compressed_model.pt2 +dp --pt-expt freeze -c model.ckpt.pt -o compressed_model --compress ``` -The two archives are not interchangeable. `frozen_model.pt2` carries the plain -graph lower and is the uncompressed intermediate; `compressed_model.pt2` carries -the compact canonical graph lower and is what you deploy. Compression selects -that lower on its own, so it takes no lower-kind option of its own. +This writes `compressed_model.pt2` directly using the default compression +settings, without compiling an intermediate uncompressed archive. Compression +selects the export schema appropriate to the model and target device. + +A branch of a multi-task model is frozen by naming it with `--head`: + +```bash +dp --pt-expt freeze -c model.ckpt.pt -o compressed_branch --head CHOSEN_BRANCH --compress +``` + +Where the branches share the fitting network, the case embedding that +distinguishes them is a constant of the branch, and the export folds it into +the bias of the first fitting layer, so the compressed branch takes the export +schema and the fused kernels of a single-task model. The fold is not possible +where the descriptor and the case embedding together are as wide as the first +fitting layer, or half as wide, because the residual connection of that layer +then carries the case embedding; such a branch keeps it and exports to the +graph form, which the export reports with a warning. + +On a CUDA device, compression also measures which kernel evaluates the fitting +network fastest on that GPU at each system size and stores the choices in the +archive. A kernel other than the default is chosen only where the measurement +finds it faster by more than 1% on a lightly loaded GPU and, under sustained +load, at most 1% slower both on a system of one element and on a system of up +to three elements. The measurement adds a few seconds to several minutes to a +compression, depending on the size of the model and on the GPU. All choices +give bit-identical results, so they affect only the speed. On a GPU of another +model, the archive uses the default choices instead; compress on the GPU model +that will run the simulations. + +To compress an existing frozen archive or choose the spline spacing, use the +compression command: + +```bash +dp --pt-expt compress -i frozen_model.pt2 -o compressed_model.pt2 --step 0.01 +``` + +For ASE, `DP(model="model.pt", freeze=True)` automatically compresses supported +DPA4C checkpoints and exports a single-process `.pt2`. If compression is not +supported, ASE warns and retains the ordinary freeze path. See the +[ASE interface](../third-party/ase.md). Only `-s, --step` applies to DPA4C; it sets the uniform spline spacing in Å, and a smaller value means a finer table and a larger model. The `--extrapolate`, @@ -359,14 +395,15 @@ computed from data; DPA4C needs none of them and ignores them. Compression requires: -- the PyTorch Exportable backend on CUDA; +- the PyTorch Exportable backend; - `precision: "float32"`; - `channels`, `lmax` and `radial_modes` inside the compiled sets listed above; - an empty {ref}`exclude_types `, since the fused kernel has no type-exclusion branch. -`dp --pt-expt compress` reports an explicit error when any of these is not met. +`dp --pt-expt compress` and `freeze --compress` report an explicit error when +any of these is not met. A model that excludes type pairs still trains and runs; deploy it as the uncompressed graph archive. @@ -404,6 +441,14 @@ Run under Kokkos with one GPU: lmp -k on g 1 -sf kk -in in.lammps ``` +On a single MPI rank, the Kokkos styles replay the evaluation of the model as +a CUDA graph from step to step. A step that also computes the per-atom virial, +which `compute centroid/stress/atom` requests (the Kokkos styles provide the +nine-component per-atom virial only), evaluates more outputs than the steps +around it: that step and the next one run without the replay, which resumes +one step later. Requesting the per-atom virial at a coarse interval keeps this +cost small. + ### CPU hosts DPA4C has a second set of hand-written operators for the CPU, so a compressed @@ -475,6 +520,18 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 mpirun -np 4 lmp -in in.lammps Use a non-zero neighbor skin, for example `neighbor 2.0 bin`, to keep per-step GPU memory stable; a zero skin rebuilds the neighbor list every step. +Use an MPI library that can access GPU memory (CUDA-aware MPI); +[Install LAMMPS](../install/install-lammps.md#mpi-for-runs-on-several-gpus) +describes the setup. The KOKKOS package of LAMMPS then exchanges atom data +between ranks directly in GPU memory. When LAMMPS detects an MPI library without +this support, it prints a warning, turns GPU-aware communication off and copies +the whole position array through host memory at every step; +`pair_style deepmd/kk` and `pair_style dpa4spin/kk` then send the forces of +ghost atoms to the ranks that own them themselves, staging only those values in +host memory. On three RTX PRO 6000 GPUs with one million carbon atoms, a step of +the Nano model took 6.8 ms with a CUDA-aware Open MPI and 11.3 ms with an MPI +library without CUDA support, against 15.2 ms on one GPU. + ## Inference settings Inference behavior is controlled by environment variables read when the model is @@ -589,12 +646,15 @@ from a broken data pipeline and reports a vanishing magnetic force. ### Running a spin model in LAMMPS Freeze and compress exactly as for any other DPA4C model, then select a spin -pair style from the table under -[Choosing a pair style](#choosing-a-pair-style). Both spin styles require -`atom_style spin`: +pair style from the table under [Choosing a pair style](#choosing-a-pair-style). +Both spin styles require `atom_style spin` and, on a single MPI rank, an atom +map; the +[LAMMPS command reference](../third-party/lammps-command.md#pair_style-dpa4spin) +describes their keywords. A minimal input reads: ```lammps atom_style spin +atom_modify map yes pair_style dpa4spin compressed_model.pt2 pair_coeff * * Ni O ``` @@ -604,10 +664,27 @@ already performs, and the magnetic force is reduced back onto owning atoms alongside the conservative force, so domain decomposition needs no additional exchange. -`min_style spin` reads the magnetic force from the pair style and relaxes the -moment directions. Spin dynamics through `fix nve/spin` requires a LAMMPS build -whose fix recognizes this pair style, because the stock fix accumulates the -magnetic force only from pair styles matching its own name pattern. +By default the pair styles store the precession magnetic force that the LAMMPS +SPIN package expects, the negative energy gradient `-dE/dS` scaled by the +moment magnitude over `hbar`. `min_style spin` reads this force and relaxes the +moment directions. The `fix nve/spin` of a stock LAMMPS build does not use this +force: the fix evaluates the magnetic force itself, through the pair styles +whose names begin with `spin`, and `dpa4spin` and `dpa4spin/kk` are not among +them. With no such pair style, the fix runs silently on a single MPI rank and +stops with a sectoring error on several. Spin dynamics therefore needs a LAMMPS +build whose integrator reads the magnetic force of these pair styles. + +A Newtonian integrator of the Cartesian spin vector needs `-dE/dS` itself +rather than the precession force. The keyword `magnetic_force cartesian` +selects `-dE/dS` and requires `units metal`: + +```lammps +pair_style dpa4spin/kk compressed_model.pt2 magnetic_force cartesian +``` + +The pair styles declare the selected convention through the scalar integer +`Pair::extract("fm_cartesian_force", dim)`, zero for precession and one for +Cartesian, so such an integrator can reject a pair style left at the default. A worked example, a rocksalt NiO cell in its type-II antiferromagnetic order, is provided in `examples/spin/dpa4c/lmp/`. diff --git a/doc/model/index.rst b/doc/model/index.rst index c6db98d169..be72573447 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -13,6 +13,7 @@ Model dpa3 dpa4 dpa4c + long-range unimol train-hybrid sel @@ -22,6 +23,7 @@ Model train-fitting-tensor train-fitting-dos train-fitting-property + train-fitting-density train-se-e2-a-tebd train-se-a-mask train-se-e3-tebd diff --git a/doc/model/long-range.md b/doc/model/long-range.md new file mode 100644 index 0000000000..8797277cea --- /dev/null +++ b/doc/model/long-range.md @@ -0,0 +1,256 @@ +# Long-range latent charges {{ pytorch_icon }} + +> [!NOTE] +> **Supported backends**: PyTorch Exportable {{ pytorch_icon }} (`dp --pt-expt`) + +A short-range model sees each atom only through the neighbors within its +cutoff. Interactions that decay slowly with distance, such as those between +charged groups of a molecule, therefore escape it. The long-range module adds +such an interaction to the energy of a standard model. A second network reads +the same descriptor as the energy fitting and predicts per-atom _latent +charges_, and every pair of atoms of a system interacts through those charges: + +```math +E = \sum_i E_i^{\text{SR}} + k_e \sum_{i [!WARNING] +> The sum covers complete systems. Without `periodic: true`, a periodic cell +> is rejected. The cost of the non-periodic sum grows with the square of the +> number of atoms of a frame. + +## Configuration + +The module is a `long_range` section of a standard energy model, beside the +descriptor and the fitting. It is currently available with the +[DPA4C](./dpa4c.md) descriptor: + +```json +"model": { + "type_map": ["C", "H", "O"], + "descriptor": { + "type": "dpa4c", + "rcut": 6.0 + }, + "fitting_net": { + "neuron": [128, 128, 128] + }, + "long_range": { + "type": "latent_charge", + "channels": 1, + "kernel": { + "type": "les", + "alpha": 1.0 + }, + "charge_constraint": { + "source": "neutral" + } + } +} +``` + +The charge network, +{ref}`charge_net `, +takes the frame and atomic parameters of the energy fitting and follows its +activation function, skip connection and precision unless they are set; its +seed defaults to one derived from the seed of the fitting. +Its width is set by +{ref}`neuron `. + +### Kernels + +{ref}`kernel ` selects the +interaction of the charges, `les` by default. Both kernels are trained with +the model unless their `trainable` option is false, and their widths stay +positive throughout training. + +- `les`: $K(r) = \mathrm{erf}(\alpha r)/r$, the interaction of two Gaussian + charge clouds. It approaches $1/r$ beyond a few $1/\alpha$ and stays finite, + $2\alpha/\sqrt{\pi}$, at contact. The inverse width $\alpha$ is set by + {ref}`alpha `. +- `sog`: $K(r) = \sum_m a_m \exp(-r^2/2\sigma_m^2)$, a sum of Gaussians whose + amplitudes $a_m$ may take either sign. The widths are given by + {ref}`width ` + or default to a geometric sequence, and the amplitudes by + {ref}`amplitude `, + which explicit widths require. The default amplitudes, + $a_m = \sqrt{2/\pi}\,\ln b/\sigma_m$ for the ratio $b$ of consecutive + widths, discretize $1/r$: with the default widths the kernel stays finite, + $0.46$ Å⁻¹, at contact, and its radial force between fixed charges stays + within 0.4% of that of $1/r$ from 6.5 Å to 150 Å. It differs from $1/r$ + there by a nearly constant offset $C \approx 1.3 \times 10^{-3}$ Å⁻¹, which + exerts no force between fixed charges. The charges respond to the positions, + however, and through them the offset adds + $-k_e C \sum_c (Q_c^2 - \sum_i (q_i^c)^2)/2$ to the energy of a frame, with + $Q_c$ the total charge of channel $c$, which the response of the charges + turns into a force. Below about 5 Å the kernel falls short of $1/r$ by more + than $C$, by 0.12 Å⁻¹ at 2 Å. + +Several charge channels, set by +{ref}`channels `, let +the model represent interactions that one charge per atom cannot. + +### Periodic systems + +With +{ref}`periodic ` enabled, +a frame given a cell is evaluated over the images of the cell. The sum +converges in reciprocal space, where the Fourier transform of the kernel, + +- `les`: $\hat K(\mathbf{k}) = (4\pi/k^2)\, e^{-k^2/4\alpha^2}$, or +- `sog`: $\hat K(\mathbf{k}) = \sum_m a_m (2\pi)^{3/2} \sigma_m^3\, e^{-\sigma_m^2 k^2/2}$, + which with the default amplitudes approaches the Coulomb transform + $4\pi/k^2$, + +decays like a Gaussian. Over the wave numbers $\mathbf{k} = 2\pi B^{-1} +\mathbf{n}$, $\mathbf{n} \in \mathbb{Z}^3 \setminus \{\mathbf{0}\}$, of a cell +$B$ whose rows are the lattice vectors, + +```math +E = k_e \left[\frac{1}{2V}\sum_{\mathbf{k}\neq\mathbf{0}} \sum_{c=1}^{C} +\hat K(\mathbf{k})\, |\hat\rho_c(\mathbf{k})|^2 +- \frac{1}{2} K(0) \sum_{i,c} (q_i^c)^2 \right], +\qquad \hat\rho_c(\mathbf{k}) = \sum_i q_i^c e^{i\mathbf{k}\cdot\mathbf{r}_i}, +``` + +with the cell volume $V$ and the kernel value at contact $K(0)$, whose term +removes the interaction of every charge with itself. The +$\mathbf{k}=\mathbf{0}$ mode is dropped, the tin-foil convention, which places +a uniform neutralizing background under a channel with a net charge. The +wave-number cutoff $k_{max}$ is where the weight of the slowest-decaying mode +falls below +{ref}`kspace_tolerance `, +so the truncation converges exponentially; it follows the kernel as it trains, +unless +{ref}`decouple_k_grid ` +keeps the cutoff of the initial kernel. + +{ref}`k_grid_mode ` +selects how training evaluates the sum: + +- `pme`, the default: by the smooth particle-mesh Ewald method. The charges + are spread with B-splines of order + {ref}`pme_order ` to a + mesh that resolves every cell of the batch at + {ref}`pme_spacing `, + whose transform gives the sum; the cost grows nearly linearly with the + number of atoms. With the default order 8 and spacing $0.75\pi/k_{max}$, the + energy and the forces stay within a few parts in a million of the sum over + the wave numbers. +- `ragged`: over the wave numbers of every frame within $k_{max}$, listed for + every batch; their number grows with the volume of the cell. +- `fixed`: over one integer grid shared by the frames of a batch, sized by the + largest cell of the data statistics or by + {ref}`k_grid_bounds `; + a frame whose cell needs a wave number beyond the grid is truncated to it, + with a warning. + +A frozen model evaluates every periodic frame on the mesh of its cell, +whatever its training mode, so a model trained in another mode can differ +from its frozen form by the accuracy of the mesh. Forces and virials hold the +complete periodic term, and a model with `periodic: true` still evaluates +non-periodic frames by the direct sum. In a `linear_ener` composition, a +sub-model whose module serves periodic systems takes no other sub-model with +a long-range module beside it. + +### Total charge + +{ref}`charge_constraint ` +shifts one channel uniformly over the atoms of every frame so that it sums to +the total charge of the frame. The total charge is zero (`neutral`), the +charge of the frame charge/spin condition of the descriptor (`charge_spin`), +or a column of the frame parameters (`fparam`); a frame without the input +takes its default. The constrained channel then carries the total charge, +while the other channels remain free. Without the section, no channel is +constrained. + +### Exclusions + +Pairs of the types in `pair_exclude_types` see each other neither in the +descriptor nor in the long-range interaction, while the `exclude_types` of +the descriptor act on the descriptor alone. In a `linear_ener` composition, +`pair_exclude_types` belongs to the composition, whose sub-models share its +neighbor graph and pair topology. The reciprocal-space sum of a periodic +system is global and excludes no pair, so `pair_exclude_types` is not +supported with a module that serves periodic systems, nor +`atom_exclude_types` with any long-range module. + +## Training and inference + +Training uses the usual energy loss; see +[training energy models](train-energy.md). With `enable_compile`, the +long-range interaction is part of the compiled model: the pairs, the wave +numbers or the mesh of each batch are built outside the compiled region and +enter it as inputs of dynamic size. Compiled training of a periodic model in +the `pme` mode on the CPU requires PyTorch 2.13 or later. + +`dp --pt-expt freeze` exports the model, long-range module included, to +`.pt2` or `.pte`, and `dp --pt-expt compress` compresses its descriptor into a +`.pt2` archive while the networks of the long-range module stay as they are. +Freezing a model with `periodic: true` to `.pt2` requires PyTorch 2.12 or +later, except for a model that the dedicated operator described below +evaluates, whose mesh the operator transforms. +A trained checkpoint (`.pt`) and the frozen model are evaluated by `dp test` +and the Python inference interface; the C++ interface and LAMMPS, with +`pair_style deepmd` and `pair_style deepmd/kk`, read the `.pt2` archive. +Since the sum covers the complete system, LAMMPS runs a model whose module +serves no periodic system with non-periodic boundaries (`f`, `s` or `m`) in +every dimension, and a model with `periodic: true` with boundaries periodic in +every dimension or in none. Several models given to `pair_style deepmd` for +the model deviation evaluate a non-periodic system together only when every +one of them carries a long-range module. A model that the dedicated operator described +below evaluates, a compressed DPA4C model frozen for a CUDA device, also runs +on several MPI ranks: every rank gathers the positions and charges of the +atoms of all ranks over MPI and evaluates its own atoms against all of them, +so that the ranks share the cost of the direct sum, while each evaluates the +mesh of a periodic system whole. This requires DeePMD-kit built with MPI. Any +other model, like an uncompressed one, runs on a single rank; +`pair_style deepmd/kk` evaluates an uncompressed model on non-periodic systems +only. + +Compressing a descriptor conditioned on `charge_spin` specializes the model +to one charge state at a time, and a constraint that reads its total charge +from `charge_spin` follows the same state. The compressed archive serves any +charge state: the Python interface applies the state of each call, and LAMMPS +the one given by the `charge_spin` keyword of `pair_style deepmd`. + +A compressed DPA4C model frozen for a CUDA device evaluates the long-range +term with a dedicated operator, which reads the positions of the atoms and +the cell. It enumerates the pairs of a non-periodic system itself, so that its +memory grows linearly with the number of atoms and its cost with their +square, and evaluates a periodic system on the mesh of its cell. The operator +serves the `les` kernel and the `sog` kernel of up to 32 Gaussians, with up to +8 charge channels, B-splines of order up to 12, a charge network without +frame or atomic parameters, and no `pair_exclude_types`. A model the +operator does not serve, like an uncompressed one, receives the pairs or the +mesh of each frame from the interface, so the memory of a non-periodic +evaluation grows with the square of the number of atoms. The operator +evaluates the long-range term together with the descriptor and both networks, +which is faster, when the module has one charge channel and a charge network +in `float32` with `tanh` or `silu` hidden layers whose widths are multiples of +four and without `resnet_dt`; the defaults give such a module in a compressed +model. + +The following are not supported for a model with a long-range module: +multi-rank (domain-decomposed) inference of a model that the dedicated +operator does not evaluate, such as an uncompressed one, spin models, the +tabulated ZBL potential, and short-range bridging. diff --git a/doc/model/train-energy-spin.md b/doc/model/train-energy-spin.md index d06bf951bc..17a7701785 100644 --- a/doc/model/train-energy-spin.md +++ b/doc/model/train-energy-spin.md @@ -12,7 +12,9 @@ keeping other sections the same as the normal energy model's input script. > > - In the TensorFlow backend, the `se_e2_a` descriptor will treat those atom types with spin as new (virtual) types, > and duplicate their corresponding selected numbers of neighbors ({ref}`sel `) from their real atom types. -> - In the PyTorch-TorchScript and PyTorch-Exportable backends with the default `deepspin` (virtual-atom) scheme, 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, descriptors +> using the `deepspin` (virtual-atom) scheme add virtual types for all atom +> types, with or without spin. > 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, @@ -80,11 +82,10 @@ 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). +The `se_e2_a`, `dpa1`, `dpa2`, and `dpa3` descriptors use the `deepspin` +(virtual-atom) scheme. DPA4 uses the `native` scheme +({ref}`model.spin.scheme `), which passes the per-atom spin vector +directly into the descriptor; see the [DPA4 spin documentation](dpa4.md#spin). > [!NOTE] > It should be noted that the spin models in PyTorch-TorchScript/PyTorch-Exportable/DP are capable of addressing scenarios where the spin approaches zero diff --git a/doc/model/train-energy.md b/doc/model/train-energy.md index 83bf693866..92ede6e028 100644 --- a/doc/model/train-energy.md +++ b/doc/model/train-energy.md @@ -151,13 +151,15 @@ limit of the model is pinned. Two model options provide this together: `omat24` is one of the bundled tables of isolated-atom energies in eV: -| Name | Reference calculation | Elements | -| -------- | --------------------- | ----------------------------------------- | -| `mptraj` | MPtraj | 89, H to Pu without Po, At, Rn, Fr and Ra | -| `omat24` | OMat24 | 89, H to Pu without Po, At, Rn, Fr and Ra | -| `omol25` | OMol25, neutral atoms | 83, H to Bi | -| `omc25` | OMC25 | 94, H to Bk without Tb, Am and Cm | -| `odac25` | ODAC25 | 94, H to Pu | +| Name | Reference calculation | Elements | +| ---------------- | --------------------- | ----------------------------------------- | +| `domain_cluster` | Domain_cluster | 9: Al, Si, Ni, Cu, Ru, Pd, Ag, Pt, Au | +| `openmd` | OpenMD, ABACUS PBE PW | 83, H to Bi | +| `mptraj` | MPtraj | 89, H to Pu without Po, At, Rn, Fr and Ra | +| `omat24` | OMat24 | 89, H to Pu without Po, At, Rn, Fr and Ra | +| `omol25` | OMol25, neutral atoms | 83, H to Bi | +| `omc25` | OMC25 | 94, H to Bk without Tb, Am and Cm | +| `odac25` | ODAC25 | 94, H to Pu | The `mptraj` table uses the final VASP `free energy TOTEN`, in eV. @@ -166,6 +168,24 @@ reference energies of the UMA training tasks published with fairchem (`configs/uma/training_release/element_refs/iso_atom_elem_refs.yaml`, MIT license). +The `openmd` table contains the OpenMD ABACUS PBE plane-wave atomic +reference energies from `ABACUS_PBE_PW_atom_E0.json` (H to Bi), in eV. +Use `"preset_out_bias": {"energy": "openmd"}` to select it. +These are ABACUS 3.9.0 `!FINAL_ETOT_IS` values calculated with the OpenMD +UPF pseudopotentials, a 100 Ry cutoff, a 30 Angstrom cubic cell, Gamma-point +sampling, `nspin=1`, no SOC, and Gaussian smearing of 0.01 Ry. +They are finite-smearing, non-spin-polarized references, rather than +spin-polarized free-atom ground-state energies. The La value +(-950.8377156904378 eV) is the accepted final-step energy after 120 SCF +iterations; its final density residual was 1.2043e-7 against a target of +1e-8. The other 82 elements met that SCF threshold. + +The `domain_cluster` table contains the supplied Domain_cluster atomic energy +reference biases in eV for Al, Si, Ni, Cu, Ru, Pd, Ag, Pt, and Au. Select it +with `"preset_out_bias": {"energy": "domain_cluster"}`. Elements outside these +nine entries are unassigned; every element present in the training data must +have a reference value. + Each table is on the energy scale of its own reference calculation (the OMat24 table gives H = -1.117 eV, the OMol25 table H = -13.446 eV), so a model takes the table of the calculation that produced its training data. A @@ -211,9 +231,13 @@ the autograd route instead. `vacuum_ref` is available for DPA4/SeZM on the PyTorch backend and for the graph-native models (DPA1, DPA2, DPA4 and DPA4C) on the pt_expt backend, whose -flat node axis carries the reference atoms. Fine-tuning and `dp change-bias` -keep the reference; a `preset_out_bias` given to them fixes the bias again -without statistics. +flat node axis carries the reference atoms. Fine-tuning keeps the reference +and applies its configured `preset_out_bias` without fitting assigned outputs. +Standalone `dp change-bias` also keeps the reference, but fits residuals from +the stored bias without reapplying checkpoint presets. Non-finite stored bias +cells covered by a finite checkpoint preset are recovered before residual +fitting; finite adjusted values are preserved. To set an energy bias explicitly +without statistics, use `dp change-bias --bias-value ...`. ## Loss diff --git a/doc/model/train-fitting-density.md b/doc/model/train-fitting-density.md new file mode 100644 index 0000000000..10fd86a2f1 --- /dev/null +++ b/doc/model/train-fitting-density.md @@ -0,0 +1,92 @@ +# Fit charge density on grid points {{ pytorch_icon }} + +> [!NOTE] +> **Supported backends**: PyTorch-TorchScript {{ pytorch_icon }} +> +> Training and evaluation are supported in Python (`dp train`, `dp test`, +> and the `DeepEval`/`DeepDensity` interfaces). C++/LAMMPS inference is +> **not** supported yet: the model's `forward_lower` interface is not +> implemented for density models. See the multi-backend tracking issue for +> the planned dpmodel implementation and backend wrappers. + +Here we present an API to the grid density model, which predicts the charge density on a set of grid points for a given atomic configuration. + +A complete example can be found in + +```bash +$deepmd_source_dir/examples/density/ +``` + +with DPA-2 and DPA-3 configurations, a small QM9 dataset, and a README describing the data format in detail. Note that **the data provided along with the examples are of limited amount, and should not be used to train a production model.** + +## Data format + +The training/validation data follows the standard `deepmd/npy` format, with two additional files in each `set.000/` directory: + +- `grid.npy` of shape `[nframes, ngrid, 3]`: the coordinates of the grid points; +- `density.npy` of shape `[nframes, ngrid, 1]`: the charge density labels on the grid points. + +The number of grid points `ngrid` may differ from the number of atoms. The last entry of `type_map` is reserved as a virtual "grid point type": internally, grid points are assigned this type when building the grid-to-atom neighbor list, so `type_map` must contain one more entry than the real element types. For descriptors with a per-type `sel` list (e.g. `se_e2_a`), set its last entry to 0, since grid points are only centers, never neighbors. + +Every frame in one system must use the same positive `ngrid`; empty grids are +rejected. Python inference evaluates large grids in chunks of 4096 points by +default, which bounds the temporary directional-neighbor matrix. Pass +`grid_batch_size=None` to `DeepDensity.eval` only when evaluating the complete +grid in one call is known to fit in memory. + +It is recommended to set a positive `env_protection` in the descriptor configuration (e.g. `"env_protection": 0.1`). Grid points may legitimately coincide with atoms, and the default `env_protection = 0.0` would let the `1/r` terms in the environment matrix produce `NaN` densities. If `env_protection` is not set for a density model, it is automatically set to `1e-6` with a warning. + +## The fitting network + +The {ref}`fitting_net ` section tells DP which fitting net to use. + +The JSON of `density` type should be provided like + +```json + "fitting_net" : { + "type": "density", + "neuron": [120,120,120], + "resnet_dt": true, + "seed": 1 + }, +``` + +- `type` specifies which type of fitting net should be used. It should be `density`. +- The rest arguments have the same meaning as they do in `ener` mode. `numb_aparam` is not supported, because the fitting net consumes the descriptor rows of the grid points, which have no per-atom parameters. + +## Loss + +The grid density model is trained with the `grid_density` loss, which minimises the squared error between the predicted and the reference densities on all grid points, consistent with the other losses in the package: + +```json + "loss" : { + "type": "grid_density", + "start_pref_d": 1.0, + "limit_pref_d": 1.0 + }, +``` + +- {ref}`type ` should be written as `grid_density`. +- `start_pref_d` and `limit_pref_d` specify the weight of the density loss at the start and at the end of the training. If both are set to 0, the density term is skipped (and the graph is kept connected). When either prefactor is nonzero, `density.npy` is **required** in every training and validation system; when both are zero, the file is optional. + +The loss gives every grid point equal weight. Non-uniform quadrature weights, +integrated-charge constraints, and output non-negativity are not imposed by +this fitting type; datasets that need those properties must account for them +in their sampling or preprocessing. + +## Evaluation + +`dp test` evaluates the model on a system containing `grid.npy` and `density.npy` and reports the MAE and RMSE of the density. For programmatic access, `DeepEval` dispatches density models to `DeepDensity`, whose `eval` takes the grid coordinates and returns the density on the grid points: + +```python +from deepmd.infer import DeepEval + +dp = DeepEval("model.pth") +density = dp.eval( + coord, + box, + atype, + grid=grid, + grid_batch_size=4096, +) # (nframes, ngrid) +``` diff --git a/doc/third-party/ase.md b/doc/third-party/ase.md index f4c350fca1..79595a406d 100644 --- a/doc/third-party/ase.md +++ b/doc/third-party/ase.md @@ -60,6 +60,81 @@ print(water.get_forces()) :::: +Neighbor construction can be configured independently of checkpoint freezing: + +```python +calculator = DP( + model="model.pt", + neighbor_graph_method="auto", + neighbor_skin=1.0, +) +``` + +For PT-expt graph models, `neighbor_graph_method` accepts `auto`, `dense`, `ase`, +`cell`, `vesin`, and `nv`. On a GPU, `auto` prefers Vesin for the single frame +evaluated by ASE; batched `DeepPot` inference prefers NV when available. CPU +graph inference retains the native `cell` search. PT SeZM/DPA4 checkpoints and +their edge-vector exports support `auto`, `vesin`, and `nv`. Their `auto` policy +prefers Vesin for a single frame and NV for multiple frames on both CPU and CUDA; +explicit `nv` also supports both devices. Models with the traditional +extended-coordinate input use `nlist_backend` with the same inference policy. + +The default `neighbor_skin=1.0`, in angstrom, enables candidate reuse for Vesin +searches across all these input formats and requires Vesin and its torch +bindings at version 0.6 or later (PyTorch 2.3 or later). +Older PyTorch targets and macOS Intel installations retain native neighbor +construction without requiring these bindings. +If neither optional search library is available, edge-vector models use the +native dense carry-all graph search. This slower fallback preserves every +cutoff neighbor, including periodic images, independently of `sel`. A runtime +warning reports its O(N²) cost and installation alternatives; large training or +inference systems should use a cell-list backend to avoid excessive memory use. +Candidates extend to the model cutoff plus the skin and are rebuilt when an +atom moves more than half the skin, or when the cell or atom count changes. +Each evaluation still updates geometry and selects neighbors at the model's +exact cutoff. Set `neighbor_skin=0.0` to disable candidate reuse. Other search +backends retain their own rebuild behavior. Element-to-model type mappings +are reused until atomic numbers, atom order, or the supplied `type_dict` changes. + +For PT DPA4/SeZM checkpoints and graph-eligible PT-expt checkpoints, enable +automatic freezing for production inference: + +```python +from deepmd.calculator import DP + +calculator = DP(model="model.pt", freeze=True) +``` + +Each calculator construction with `freeze=True` rebuilds the artifact using +the selected checkpoint head and the backend's normal freeze policy. The +calculator loads a private artifact retained for its evaluator's lifetime and +atomically publishes a separate copy as `model-ase_single.pt2` beside the +checkpoint. Concurrent processes retain their own selected heads; the last +publication replaces the shared copy. This requires one extra artifact copy +per live evaluator. Concurrent exports must use separate processes because +the backend compiler configuration is not thread-safe. + +The export omits the multi-rank communication artifact +because ASE inference runs in one process. For CUDA targets, DPA4 defaults to +`DP_CUDA_INFER=1` and `DP_TRITON_INFER=2`, with explicit environment settings +taking precedence; DPA4C keeps its required CUDA level of 2. + +DPA4C checkpoints also enable compression through the same export path as +`dp --pt-expt freeze --compress`. Supported models use the dedicated CPU or +CUDA compression kernels. If the configuration cannot be compressed, such as +a `float64` descriptor or one with type exclusions, ASE logs the reason and +continues with an ordinary frozen model. Other model families keep their +uncompressed freeze behavior. Explicit `freeze --compress` reports an error +for an unsupported compression configuration. + +The default `freeze=False` loads checkpoints directly and logs a warning for +these models: inference can be **substantially slower**. This mode is intended +for debugging and short runs that avoid compilation overhead. To reuse an +export without compiling again, pass the `.pt2` file directly after concurrent +exports have finished. Already frozen models, other PT models, and PT-expt +models without graph eligibility ignore `freeze` and keep their normal loading +path. + Optimization is also available: ```python diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md index 19bed4f72c..5142a8e4ce 100644 --- a/doc/third-party/lammps-command.md +++ b/doc/third-party/lammps-command.md @@ -146,7 +146,7 @@ pair_style deepspin models ... keyword value ... - keyword = _out_file_ or _out_freq_ or _fparam_ or _fparam_from_compute_ or _aparam_from_compute_ or _charge_spin_ or _atomic_ or _relative_ or _aparam_ or _ttm_ > [!NOTE] -> Please note that the virial and atomic virial are not currently supported in spin models. +> Please note that the `deepspin` pair style does not provide the six-component per-atom virial of `compute stress/atom`; use `compute centroid/stress/atom` for the nine-component per-atom virial. 
     out_file value = filename
@@ -203,12 +203,58 @@ The unit follows [LAMMPS units](#units) and the [scale factor](https://docs.lamm
 Other settings and output for this pair style is the same as `deepmd` pair style, please see the detailed description [above](#pair_style-deepmd).
 
 > [!NOTE]
-> Please note that the virial and atomic virial are not currently supported in spin models.
+> Please note that the `deepspin` pair style does not provide the six-component per-atom virial of `compute stress/atom`; use `compute centroid/stress/atom` for the nine-component per-atom virial.
 
 ### Restrictions
 
 - The `deepspin` pair style is provided in the USER-DEEPMD package, which is compiled from the DeePMD-kit, visit the [DeePMD-kit website](https://github.com/deepmodeling/deepmd-kit) for more information.
 
+## pair_style `dpa4spin`
+
+The DeePMD-kit package provides the pair_style `dpa4spin` for native-spin [DPA4C](../model/dpa4c.md) models, which take the magnetic moment of each atom as an input and create no virtual atoms. The Kokkos variant `dpa4spin/kk` keeps the graph and the moments in device memory. For an example LAMMPS input one may check [the example input file for pair_style `dpa4spin`](../../examples/spin/dpa4c/lmp/in.lammps).
+
+```lammps
+pair_style dpa4spin model keyword value ...
+```
+
+- dpa4spin = style of this pair_style
+- model = frozen native-spin model to compute the interaction
+- keyword = _charge_spin_ or _magnetic_force_
+
+
+    charge_spin value = parameters
+        parameters = the dim_chg_spin charge/spin values that a model trained with a charge/spin embedding takes for the whole run. If omitted, the model's stored default_chg_spin is used.
+    magnetic_force value = precession or cartesian
+        precession = store the precession magnetic force that the LAMMPS SPIN package expects (default).
+        cartesian = store the negative energy gradient -dE/dS with respect to the moment; requires units metal.
+

+
+### Examples
+
+```lammps
+atom_style spin
+atom_modify map yes
+pair_style dpa4spin frozen_model.pt2
+pair_coeff * * Ni O
+
+pair_style dpa4spin frozen_model.pt2 charge_spin 1.0 2.0
+pair_style dpa4spin/kk compressed_model.pt2 magnetic_force cartesian
+```
+
+### Description
+
+Evaluate the interaction of a system with magnetic moments by a native-spin DPA4C model, which returns the energy, the force and the magnetic force of every atom. The pair style computes the global virial and, through `compute centroid/stress/atom`, the nine-component per-atom virial; `compute stress/atom` stops with an error, because the pair style provides no six-component per-atom virial.
+
+By default the pair style stores the precession magnetic force that the LAMMPS SPIN package expects, the negative energy gradient `-dE/dS` scaled by the moment magnitude over `hbar`. A Newtonian integrator of the Cartesian spin vector needs `-dE/dS` itself rather than the precession force; the keyword `magnetic_force cartesian` selects `-dE/dS`. The pair style declares the selected convention through the scalar integer `Pair::extract("fm_cartesian_force", dim)`, zero for precession and one for Cartesian, so that such an integrator can reject a pair style left at the default.
+
+### Restrictions
+
+- The `dpa4spin` pair style requires `atom_style spin`, and on a single MPI rank an atom map (`atom_modify map yes`).
+- `dpa4spin/kk` requires a compressed archive and a LAMMPS build with the KOKKOS package on the GPU backend.
+- In [plugin mode](../install/install-lammps.md#install-lammps-plugin-mode), `dpa4spin/kk` requires a DeePMD-kit plugin built with `DEEPMD_LAMMPS_KOKKOS=ON`.
+- The `fix nve/spin` of a stock LAMMPS build does not use the magnetic force of this pair style: the fix evaluates the magnetic force itself, through the pair styles whose names begin with `spin`. With no such pair style, the fix runs silently on a single MPI rank and stops with a sectoring error on several. Spin dynamics needs a LAMMPS build whose integrator reads the magnetic force of the pair style.
+- The `dpa4spin` pair style is provided in the USER-DEEPMD package, which is compiled from the DeePMD-kit, visit the [DeePMD-kit website](https://github.com/deepmodeling/deepmd-kit) for more information.
+
 ## compute `deepmd/fparam/dedn`
 
 The `deepmd/fparam/dedn` compute evaluates a finite-difference derivative of the model energy with respect to a chosen frame parameter source.
diff --git a/doc/train/finetuning.md b/doc/train/finetuning.md
index 5995c79836..273d6bf628 100644
--- a/doc/train/finetuning.md
+++ b/doc/train/finetuning.md
@@ -74,7 +74,8 @@ If you wish to conduct fine-tuning using a randomly initialized fitting net in t
 $ dp --pt train input.json --finetune pretrained.pt --model-branch RANDOM
 ```
 
-The model section in input.json **must be the same as that in the pretrained model**.
+The model section in input.json **must be the same as that in the pretrained model**,
+apart from the native-spin initialization described below.
 If you do not know the model params in the pretrained model, you can add `--use-pretrain-script` in the fine-tuning command:
 
 ```bash
@@ -91,6 +92,26 @@ The model section will be overwritten (except the `type_map` subsection) by that
     }
 ```
 
+#### Initializing native spin from a spin-free checkpoint
+
+PyTorch DPA4/SeZM and PyTorch Exportable DPA4/DPA4C models support spin-free
+pretraining with a native-spin target. The source and target must use the same
+descriptor family and identical ordered `type_map` values. Set the target's
+`spin.scheme` to `native` and configure `use_spin`; these settings are retained
+even with `--use-pretrain-script`. The same rule applies to each selected branch
+in multi-task fine-tuning.
+
+For DPA4C with `dpa4_ener` or `sezm_ener` fitting, automatic hidden widths
+(an omitted `neuron` or zero entries) inherit the pretrained fitting widths.
+The resolved widths are stored in the target configuration so that checkpoint
+reconstruction uses the same architecture. Explicit widths are not rewritten
+and must remain compatible with weight transfer. `RANDOM` instead initializes
+the target fitting network and resolves its automatic widths from the target.
+
+The transfer retains pretrained geometric statistics and initializes new spin
+routes with a closed gate. DPA4C's per-type spin-reference magnitudes are measured
+from the fine-tuning data before adjusting the output bias.
+
 #### Fine-tuning from a multi-task pre-trained model
 
 Additionally, within the PyTorch implementation and leveraging the flexibility offered by the framework and the multi-task training process proposed in DPA2 [paper](https://doi.org/10.1038/s41524-024-01493-2),
diff --git a/doc/train/multi-task-training.md b/doc/train/multi-task-training.md
index 6ebe5b4fbb..d49f9075fc 100644
--- a/doc/train/multi-task-training.md
+++ b/doc/train/multi-task-training.md
@@ -76,7 +76,7 @@ Specifically, there are several parts that need to be modified:
 
 - (Optional) {ref}`training/model_prob `: The sampling weight settings corresponding to each `model_key`, i.e., the probability weight in the training step.
   You can specify any positive real number weight for each task. The higher the weight, the higher the probability of being sampled in each training.
-  This setting is optional, and if not set, tasks will be sampled with weights proportional to the number of systems in each task. It is only used when `num_epoch_dict` is not set.
+  This setting is optional and is only used when `num_epoch_dict` is not set. The default weights use each task's native data size: the number of NPY/HDF5 systems, or the number of retained LMDB frames on the PyTorch and PyTorch Exportable backends. Grouped inputs add these native sizes, without applying within-task group probabilities. Both PyTorch backends follow the established PyTorch rule; default task weights do not depend on batch size or nominal epoch length.
 
 - (Optional) {ref}`training/num_epoch_dict `: The number of training epochs for each model branch, specified as a dictionary mapping `model_key` to epoch values (can be fractional).
   This allows different tasks to train for different numbers of epochs, which is particularly useful for multi-task fine-tuning scenarios
diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md
index 31609c09b8..353eda9bb5 100644
--- a/doc/train/training-advanced.md
+++ b/doc/train/training-advanced.md
@@ -53,11 +53,12 @@ Other training parameters are given in the {ref}`training ` section.
 The sections {ref}`training_data ` and {ref}`validation_data ` give the training dataset and validation dataset, respectively. Taking the training dataset for example, the keys are explained below:
 
 - {ref}`systems ` provide paths of the training data systems. DeePMD-kit allows you to provide multiple systems with different numbers of atoms. This key can be a `list` or a `str`.
-  - `str`: {ref}`systems ` should be a valid path. It can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories.
+  - `str`: {ref}`systems ` should be a valid path. It can be a system directory path (containing 'type.raw'), a DeePMD HDF5 container or group, or a parent directory path to recursively search for all system subdirectories.
   - `list`: {ref}`systems ` gives a list of paths. Each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories.
+  - The PyTorch and PyTorch Exportable backends also accept LMDB paths, lists of LMDB paths, and parent directories containing LMDB datasets. An LMDB dataset is terminal: discovery does not inspect its contents as additional datasets. NPY and LMDB sources may appear in the same group. Other backends discover DeePMD NPY/HDF5 systems only.
 - At each training step, DeePMD-kit randomly picks {ref}`batch_size ` frame(s) from one of the systems. The probability of using a system is by default in proportion to the number of batches in the system. More options are available for automatically determining the probability of using systems. One can set the key {ref}`auto_prob ` to
-  - `"prob_uniform"` all systems are used with the same probability.
-  - `"prob_sys_size"` the probability of using a system is proportional to its size (number of frames).
+  - `"prob_uniform"` each retained system supplies the same long-run share of sampled frames. Batch-selection probabilities account for the number of frames returned per batch.
+  - `"prob_sys_size"` the probability of using a system is proportional to its number of batches.
   - `"prob_sys_size; sidx_0:eidx_0:w_0; sidx_1:eidx_1:w_1;..."` the `list` of systems is divided into blocks. Block `i` has systems ranging from `sidx_i` to `eidx_i`. The probability of using a system from block `i` is proportional to `w_i`. Within one block, the probability of using a system is proportional to its size.
 - An example of using `"auto_prob"` is given below. The probability of using `systems[2]` is 0.4, and the sum of the probabilities of using `systems[0]` and `systems[1]` is 0.6. If the number of frames in `systems[1]` is twice of `system[0]`, then the probability of using `system[1]` is 0.4 and that of `system[0]` is 0.2.
 
@@ -103,6 +104,59 @@ The sections {ref}`training_data ` and {ref}`validation_
     Filling batches also makes the *sampling weight* more faithful. Per-atom loss terms (force, atomic energy) pool over a batch's real labels, so a frame already counts in proportion to its atom count there. Frame-level terms (energy, virial) weigh the frames of a batch equally, so a frame counts for `1 / nframes`, and an atom budget makes `nframes` follow the atom count. `"max:N"` distorts that wherever an atom-count group is too small to fill a batch — its few frames form a short batch and each is over-weighted many times over — while `"mix:N"` packs them with their neighbours and lands much closer to the intended weighting.
 - The key {ref}`numb_batch ` in {ref}`validate_data ` gives the number of batches of model validation. Note that the batches may not be from the same system
 
+### Data groups
+
+`training_data` accepts either the dictionary shown above or a nonempty list of dictionaries. Each dictionary defines a group with its own `systems`, `batch_size`, and within-group sampling options. The optional `prob` specifies the probability of drawing a training batch from the entire group, including every dataset found recursively under its paths:
+
+```json
+"training_data": [
+    {
+        "systems": "../data/collection_a",
+        "batch_size": "auto:32",
+        "prob": 0.6
+    },
+    {
+        "systems": ["../data/collection_b", "../data/extra"],
+        "batch_size": "auto:64",
+        "prob": 0.4
+    }
+]
+```
+
+The first group's datasets together receive probability 0.6, regardless of how many datasets its path contains. The second group's datasets together receive probability 0.4. Batches do not combine frames from different groups.
+
+`prob` is an absolute probability between 0 and 1. If every group specifies it, the probabilities must sum to 1. Groups without `prob` share the remaining probability in proportion to their batch counts. Omitting `prob` everywhere therefore preserves size-based sampling; it does not give equally sized shares to differently sized groups. Within each group, `auto_prob` determines the system sampling rule; explicit `sys_probs` specify batch-selection probabilities.
+
+For a group containing exactly one LMDB, `auto_prob` addresses that LMDB's internal systems, independently of the group's `prob`:
+
+```json
+"training_data": [
+    {
+        "systems": "../data/first.lmdb",
+        "batch_size": "filter:450",
+        "auto_prob": "prob_sys_size;0:10:0.5;10:30:0.5",
+        "prob": 0.6
+    },
+    {
+        "systems": "../data/second.lmdb",
+        "batch_size": "filter:450",
+        "prob": 0.4
+    }
+]
+```
+
+If a group expands to multiple sources and includes LMDB, `auto_prob` can be `"prob_sys_size"` (the default) or `"prob_uniform"`, and `sys_probs` must be omitted. Custom block rules still require a single LMDB because their system indices are local to that dataset. To assign source weights or configure different block rules, put each LMDB in a separate group. Existing dictionary configurations, NPY `sys_probs`, and single-LMDB block rules remain supported.
+
+For `"prob_uniform"`, each NPY/HDF5 system and each nonempty internal LMDB system counts once after `filter:N` is applied. An LMDB without internal system identifiers counts as one system. The rule is independent of how systems are split among LMDB files: systems containing 100 and 900 frames each supply half the sampled frames. LMDB sampling balances the internal frame counts before applying the native batch packing; `mix:N` continues to combine systems. Source-selection weights also correct for differing batch frame counts. Uniform frame shares do not imply equal contributions to each batch or to the loss, and subsequent geometric frame rejection can change the retained training shares.
+
+For an LMDB containing `N` retained frames in `M` nonempty systems, a uniform sampling pass draws `ceil(N / M)` frames from each system. Before distributed batch padding, the pass therefore contains fewer than `N + M` frame indices, rather than scaling with the largest system. Larger systems contribute a random subset without replacement; smaller systems retain their frames and repeat them to fill their quota. A pass need not visit every frame of a larger system. Subsets change with the sampler epoch and remain reproducible from its seed. `numb_epoch` and omitted outer-group probabilities use the resulting native batch counts; an explicit outer `prob` is unchanged. Custom block rules retain their existing non-shrinking targets.
+
+This frame-share definition intentionally changes older NPY/HDF5 `prob_uniform` behavior when batch sizes differ. For example, systems returning batches of 1 and 10 frames receive batch probabilities of 10/11 and 1/11, rather than 1/2 each. Short native batches are included in the correction. Explicit `sys_probs` and outer group `prob` remain batch probabilities.
+
+`validation_data` remains a dictionary. Its `systems` accepts the same recursive paths and multiple-LMDB inputs on supported backends; `numb_btch` continues to apply to the whole validation dataset.
+
+### Mixed precision
+
 The section {ref}`mixed_precision ` specifies the mixed precision settings, which will enable the mixed precision training workflow for DeePMD-kit. The keys are explained below:
 
 - {ref}`output_prec ` precision used in the output tensors, only `float32` is supported currently.
@@ -145,6 +199,69 @@ checkpoints from the same family numbered above `N`. This prevents remnants of
 a longer earlier run from displacing the newly written checkpoint from the
 retention window.
 
+PyTorch Exportable optimizer checkpoints also retain each rank's per-task
+training-data cursor and task/system selection RNG. With the same immutable
+data at the same locations, batching settings, and number of ranks, a restart
+continues with the next LMDB or directory/HDF5 training batches rather than
+starting their iteration again. Data-source
+progress is compact and restored lazily; checkpoints do not store decoded
+frames or shuffle permutations. Statistics, compilation warmup, and learning
+curve evaluation do not consume this training progression.
+Distributed ranks exchange this compact state at checkpoint boundaries, not
+on every batch; rank-local filtering may legitimately advance their cursors
+by different amounts.
+
+Restarting with a different number of ranks still restores the model, optimizer,
+and training step, but warns and keeps newly initialized data cursors and
+selection RNG instead of remapping the saved rank-local progress. Data may be
+repeated, and exact data-order continuation is not guaranteed in this case.
+Malformed progress envelopes remain errors; by default, so do incompatible data
+settings at the same world size.
+
+PyTorch Exportable initializes rank-local data sampling from `training.seed`
+and the rank before constructing data sources, including directory and HDF5
+sources. Task selection has an independent RNG shared across ranks, so local
+data filtering cannot change the task sequence. Without a configured seed,
+the chief initializes that task RNG and broadcasts its state once at startup.
+Both RNG streams resume from checkpoints. The configured seed is also part of
+the checkpoint's data-progress identity: changing it requires an explicit reset,
+independently of the input format. Earlier experimental progress records without
+seed/task-RNG identity likewise require a reset; checkpoints without data
+progress retain their legacy restart behavior.
+
+Grouped inputs also preserve each group's child-selection RNG and every native
+source's cursor, including nested groups combining directory/HDF5 and LMDB data.
+The group topology, assigned probabilities, and uniform-sampling policy must
+remain compatible when restoring progress.
+
+For PyTorch Exportable, set `training.reset_data_progress` to `true` when
+intentionally restarting with relocated data or changed data settings. This
+explicitly keeps newly initialized data cursors and task/system selection RNG
+for **all** training tasks; model weights, optimizer, learning-rate schedule,
+and training step still resume. A warning explains that batches may repeat and
+exact data-order continuation is unavailable. The default is `false`, which
+retains strict data-source compatibility checks. Malformed progress envelopes
+remain errors even when resetting.
+
+Source paths are identity guards, not content fingerprints: equal frame counts,
+atom types, or LMDB transaction numbers do not prove that two datasets contain
+the same frames. Consequently, copying an unchanged dataset to another location
+requires an explicit data-progress reset; automatic exact continuation after
+relocation is not supported. Directory/HDF5 paths must also retain the same path
+spelling, including the HDF5 group. `reset_data_progress` has no effect on fresh
+training, `--init-model`, finetuning, frozen models, or older checkpoints without
+data progress.
+
+Directory/HDF5 batches now shuffle indices into a canonical, unshuffled set.
+For a single-set source, later passes therefore differ from older releases
+that shuffled their already shuffled cache. This also prevents `numb_copy`
+multiplicities from compounding over repeated passes.
+
+Older checkpoints without data progress retain the previous restart behavior.
+`--init-model`, finetuning, and EMA-only deployment snapshots do not restore data
+progress. This is a data-order guarantee, not a guarantee of bitwise-identical
+training or of restoring every stochastic data transform.
+
 **`--init-frz-model frozen_model.pb`**, initializes the training with an existing model that is stored in `frozen_model.pb`.
 
 **`--skip-neighbor-stat`** will skip calculating neighbor statistics if one is concerned about performance. Some features will be disabled.
diff --git a/examples/density/README.md b/examples/density/README.md
new file mode 100644
index 0000000000..3e24d9705e
--- /dev/null
+++ b/examples/density/README.md
@@ -0,0 +1,218 @@
+# Charge Density Prediction Example
+
+This example demonstrates how to train and evaluate a **charge density** model using DeePMD-kit with the PyTorch backend.
+
+The model predicts the charge density on a set of grid points (`grid`) for a given atomic configuration (`coord`, `atype`, `box`).
+
+______________________________________________________________________
+
+## Directory Structure
+
+```
+.
+├── dpa2/                     # DPA-2 descriptor example
+│   └── input.json
+├── dpa3/                     # DPA-3 descriptor example (recommended)
+│   └── input.json
+├── dataset/
+│   └── qm9/
+│       ├── C7H15NO_train/    # Training data (deepmd/npy format)
+│       └── C7H15NO_val/      # Validation data (deepmd/npy format)
+└── dptest_density_script.py  # Standalone evaluation script (alternative to dp test)
+```
+
+______________________________________________________________________
+
+## Data Format
+
+The training/validation data follows the standard **`deepmd/npy`** format, with two additional files in each `set.000/` directory:
+
+| File              | Shape                   | Description                        |
+| ----------------- | ----------------------- | ---------------------------------- |
+| `coord.npy`       | `[nframes, natoms * 3]` | Atomic coordinates                 |
+| `box.npy`         | `[nframes, 9]`          | Simulation cell vectors            |
+| `type.raw`        | `[natoms]`              | Atom type indices                  |
+| `type_map.raw`    | `[ntypes]`              | Real element types in this dataset |
+| **`grid.npy`**    | `[nframes, ngrid, 3]`   | **Grid point coordinates**         |
+| **`density.npy`** | `[nframes, ngrid, 1]`   | **Charge density labels on grids** |
+
+> **Notes:**
+>
+> - `grid.npy` and `density.npy` are required for the density model. The number of grid points (`ngrid`) must match between `grid.npy` and `density.npy`, and is allowed to differ from `natoms`.
+> - The **last entry of the model configuration's `type_map` is reserved as a virtual "grid point type"** (e.g. `X` in the example): internally, grid points are assigned this type when building the grid-to-atom neighbor list. Make sure the model `type_map` contains one more entry than the real element types. `type_map.raw` may list only the real types actually present in that dataset, as the bundled example does. For descriptors with a per-type `sel` list (e.g. `se_e2_a`), set its **last entry to 0** (grid points are only centers, never neighbors); scalar `sel` descriptors such as DPA-2/DPA-3 need no change. The descriptor input statistics for the grid type are computed from the actual `grid.npy` data during training initialization.
+> - Every frame in one system must contain the same positive number of grid points. Python inference chunks large grids into groups of 4096 points by default to limit temporary neighbor-list memory.
+> - **Always set a positive `env_protection`** in the descriptor configuration (e.g. `"env_protection": 0.1`, as in `dpa3/input.json`). Grid points may legitimately coincide with atoms, and the default `env_protection = 0.0` would let the `1/r` terms in the environment matrix produce `NaN` densities and poison the training gradients at such points.
+
+______________________________________________________________________
+
+## Training
+
+### 1. Choose a Configuration
+
+Two example configurations are provided:
+
+- **`dpa2/input.json`** — Uses the DPA-2 descriptor.
+- **`dpa3/input.json`** — Uses the DPA-3 descriptor (recommended).
+
+Key settings in `input.json`:
+
+```json
+{
+  "model": {
+    "type_map": [
+      "Li",
+      "Ni",
+      "Co",
+      "Mn",
+      "O",
+      "C",
+      "H",
+      "N",
+      "F",
+      "X"
+    ],
+    "descriptor": {
+      "type": "dpa3"
+    },
+    "fitting_net": {
+      "type": "density",
+      "neuron": [
+        240,
+        240,
+        240
+      ]
+    }
+  },
+  "loss": {
+    "type": "grid_density",
+    "start_pref_d": 1,
+    "limit_pref_d": 1
+  },
+  "training": {
+    "training_data": {
+      "systems": [
+        "../dataset/qm9/C7H15NO_train"
+      ],
+      "batch_size": "auto:128"
+    },
+    "validation_data": {
+      "systems": [
+        "../dataset/qm9/C7H15NO_val"
+      ],
+      "batch_size": 1,
+      "numb_btch": 3
+    }
+  }
+}
+```
+
+- The fitting net type is `"density"` and the loss type is `"grid_density"`.
+- The model type `"grid_density"` is selected automatically from the fitting net type.
+
+### 2. Run Training
+
+```bash
+cd dpa3   # or cd dpa2
+dp --pt train input.json
+```
+
+The training will output:
+
+- `model.ckpt-*.pt` — Model checkpoints
+- `lcurve.out` — Training/validation loss curves (`mae_d` / `rmse_d` for density)
+- `out.json` — Final training parameters
+
+### 3. Finetune from a Pretrained Model
+
+```bash
+cd dpa3
+dp --pt train input.json --finetune model.ckpt-*.pt
+```
+
+> **Note:** For density models, `change_out_bias` (the energy-bias adjustment used in standard finetuning) is **automatically skipped** because density outputs are grid-based, not atomic-based. The descriptor weights are inherited, and the fitting net adapts via normal gradient descent.
+
+### 4. Freeze the Model
+
+Export a trained checkpoint into a frozen model for inference:
+
+```bash
+cd dpa3
+dp --pt freeze
+```
+
+This reads `model.ckpt.pt` in the current directory and generates `frozen_model.pth`.
+
+______________________________________________________________________
+
+## Testing / Evaluation
+
+### Option 1: `dp test` (recommended)
+
+`dp test` supports density models natively:
+
+```bash
+cd dpa3
+dp --pt test -m frozen_model.pth -s ../dataset/qm9/C7H15NO_val -n 0
+```
+
+- `-n 0` tests **all** frames; use a positive number to test only the first N frames.
+- Add `--detail-file detail` to write per-frame `[label, prediction]` pairs to `detail.density.out.` (one file per frame, one line per grid point).
+- Add `--shuffle-test --rand-seed 42` to test on a shuffled subset.
+
+Example output:
+
+```
+# testing system : ../dataset/qm9/C7H15NO_val
+# number of test data : 30
+DENSITY MAE             : x.xxxe+00 units
+DENSITY RMSE            : x.xxxe+00 units
+```
+
+### Option 2: Standalone script
+
+`dptest_density_script.py` evaluates a model and writes **one combined** `[prediction, label]` file (`result.d.out`), which is convenient for plotting:
+
+```bash
+python dptest_density_script.py \
+    dpa3/frozen_model.pth \
+    dataset/qm9/C7H15NO_val \
+    --ratio 1.0 \
+    --output val_result.txt
+```
+
+| Argument      | Description                                                               |
+| ------------- | ------------------------------------------------------------------------- |
+| `model`       | Path to the frozen model file (`.pth`)                                    |
+| `data_dir`    | Root directory of deepmd/npy datasets (searched recursively)              |
+| `--ratio`     | Fraction of frames to randomly sample (default: `0.1`)                    |
+| `--output`    | If provided, save screen output to this file as well                      |
+| `--pred-file` | File to save paired `[prediction, label]` array (default: `result.d.out`) |
+
+> **Note:** Each individual system must use one fixed `ngrid`, but different
+> systems discovered by the script may use different grid extents. Large grids
+> are evaluated in chunks by the Python inference interface.
+
+### Option 3: Python interface
+
+```python
+import numpy as np
+from deepmd.infer import DeepEval
+
+dp = DeepEval("frozen_model.pth")  # dispatched to DeepDensity automatically
+density = dp.eval(
+    coord,
+    box,
+    atype,
+    grid=grid,
+    grid_batch_size=4096,
+)  # (nframes, ngrid)
+```
+
+______________________________________________________________________
+
+## Notes
+
+- **Backend:** PyTorch only (`--pt`). Make sure DeePMD-kit is installed with PyTorch support.
+- **Deployment scope:** Python inference only (`dp test`, `DeepEval`/`DeepDensity`, and the standalone script). C++/LAMMPS inference is not supported yet (`forward_lower` is not implemented for density models).
+- **Stat file:** `input.json` specifies `"stat_file": "./qm9_charge_density.hdf5"` for caching descriptor statistics; it is generated automatically on the first run.
+- **Checkpoint vs. frozen model:** evaluation requires a frozen model (`.pth`). Freeze a checkpoint first with `dp --pt freeze`.
diff --git a/examples/density/dataset/qm9/C7H15NO_train/set.000/box.npy b/examples/density/dataset/qm9/C7H15NO_train/set.000/box.npy
new file mode 100644
index 0000000000..37946ac645
Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_train/set.000/box.npy differ
diff --git a/examples/density/dataset/qm9/C7H15NO_train/set.000/coord.npy b/examples/density/dataset/qm9/C7H15NO_train/set.000/coord.npy
new file mode 100644
index 0000000000..7927d65163
Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_train/set.000/coord.npy differ
diff --git a/examples/density/dataset/qm9/C7H15NO_train/set.000/density.npy b/examples/density/dataset/qm9/C7H15NO_train/set.000/density.npy
new file mode 100644
index 0000000000..24289af963
Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_train/set.000/density.npy differ
diff --git a/examples/density/dataset/qm9/C7H15NO_train/set.000/grid.npy b/examples/density/dataset/qm9/C7H15NO_train/set.000/grid.npy
new file mode 100644
index 0000000000..fb95783aa9
Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_train/set.000/grid.npy differ
diff --git a/examples/density/dataset/qm9/C7H15NO_train/type.raw b/examples/density/dataset/qm9/C7H15NO_train/type.raw
new file mode 100644
index 0000000000..c46e48ed0a
--- /dev/null
+++ b/examples/density/dataset/qm9/C7H15NO_train/type.raw
@@ -0,0 +1,24 @@
+5
+5
+5
+5
+5
+5
+5
+4
+7
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
diff --git a/examples/density/dataset/qm9/C7H15NO_train/type_map.raw b/examples/density/dataset/qm9/C7H15NO_train/type_map.raw
new file mode 100644
index 0000000000..36c071bf5b
--- /dev/null
+++ b/examples/density/dataset/qm9/C7H15NO_train/type_map.raw
@@ -0,0 +1,9 @@
+Li
+Ni
+Co
+Mn
+O
+C
+H
+N
+F
diff --git a/examples/density/dataset/qm9/C7H15NO_val/set.000/box.npy b/examples/density/dataset/qm9/C7H15NO_val/set.000/box.npy
new file mode 100644
index 0000000000..05e520e22e
Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_val/set.000/box.npy differ
diff --git a/examples/density/dataset/qm9/C7H15NO_val/set.000/coord.npy b/examples/density/dataset/qm9/C7H15NO_val/set.000/coord.npy
new file mode 100644
index 0000000000..f6cea16e30
Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_val/set.000/coord.npy differ
diff --git a/examples/density/dataset/qm9/C7H15NO_val/set.000/density.npy b/examples/density/dataset/qm9/C7H15NO_val/set.000/density.npy
new file mode 100644
index 0000000000..e0c673e4e2
Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_val/set.000/density.npy differ
diff --git a/examples/density/dataset/qm9/C7H15NO_val/set.000/grid.npy b/examples/density/dataset/qm9/C7H15NO_val/set.000/grid.npy
new file mode 100644
index 0000000000..13a126f2bf
Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_val/set.000/grid.npy differ
diff --git a/examples/density/dataset/qm9/C7H15NO_val/type.raw b/examples/density/dataset/qm9/C7H15NO_val/type.raw
new file mode 100644
index 0000000000..c46e48ed0a
--- /dev/null
+++ b/examples/density/dataset/qm9/C7H15NO_val/type.raw
@@ -0,0 +1,24 @@
+5
+5
+5
+5
+5
+5
+5
+4
+7
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
+6
diff --git a/examples/density/dataset/qm9/C7H15NO_val/type_map.raw b/examples/density/dataset/qm9/C7H15NO_val/type_map.raw
new file mode 100644
index 0000000000..36c071bf5b
--- /dev/null
+++ b/examples/density/dataset/qm9/C7H15NO_val/type_map.raw
@@ -0,0 +1,9 @@
+Li
+Ni
+Co
+Mn
+O
+C
+H
+N
+F
diff --git a/examples/density/dpa2/input.json b/examples/density/dpa2/input.json
new file mode 100644
index 0000000000..c25fa42d2c
--- /dev/null
+++ b/examples/density/dpa2/input.json
@@ -0,0 +1,117 @@
+{
+  "_comment": "that's all",
+  "model": {
+    "type_map": [
+      "Li",
+      "Ni",
+      "Co",
+      "Mn",
+      "O",
+      "C",
+      "H",
+      "N",
+      "F",
+      "X"
+    ],
+    "descriptor": {
+      "type": "dpa2",
+      "repinit": {
+        "tebd_dim": 8,
+        "rcut": 6.0,
+        "rcut_smth": 0.5,
+        "nsel": 120,
+        "neuron": [
+          25,
+          50,
+          100
+        ],
+        "axis_neuron": 12,
+        "activation_function": "tanh",
+        "three_body_sel": 48,
+        "three_body_rcut": 4.0,
+        "three_body_rcut_smth": 3.5,
+        "use_three_body": true
+      },
+      "repformer": {
+        "rcut": 4.0,
+        "rcut_smth": 3.5,
+        "nsel": 48,
+        "nlayers": 12,
+        "g1_dim": 128,
+        "g2_dim": 32,
+        "attn2_hidden": 32,
+        "attn2_nhead": 4,
+        "attn1_hidden": 128,
+        "attn1_nhead": 4,
+        "axis_neuron": 4,
+        "update_h2": false,
+        "update_g1_has_conv": true,
+        "update_g1_has_grrg": true,
+        "update_g1_has_drrd": true,
+        "update_g1_has_attn": false,
+        "update_g2_has_g1g1": false,
+        "update_g2_has_attn": true,
+        "update_style": "res_residual",
+        "update_residual": 0.01,
+        "update_residual_init": "norm",
+        "attn2_has_gate": true,
+        "use_sqrt_nnei": true,
+        "g1_out_conv": true,
+        "g1_out_mlp": true
+      },
+      "precision": "float32",
+      "add_tebd_to_repinit_out": false,
+      "env_protection": 0.1
+    },
+    "fitting_net": {
+      "type": "density",
+      "neuron": [
+        240,
+        240,
+        240
+      ],
+      "resnet_dt": true,
+      "precision": "float32",
+      "seed": 1,
+      "_comment": " that's all"
+    },
+    "_comment": " that's all"
+  },
+  "learning_rate": {
+    "type": "exp",
+    "decay_steps": 5000,
+    "start_lr": 0.001,
+    "stop_lr": 3.51e-08,
+    "_comment": "that's all"
+  },
+  "loss": {
+    "type": "grid_density",
+    "start_pref_d": 1,
+    "limit_pref_d": 1,
+    "_comment": " that's all"
+  },
+  "training": {
+    "stat_file": "./dpa2.hdf5",
+    "training_data": {
+      "systems": [
+        "../dataset/qm9/C7H15NO_train"
+      ],
+      "batch_size": 1,
+      "_comment": "that's all"
+    },
+    "validation_data": {
+      "systems": [
+        "../dataset/qm9/C7H15NO_val"
+      ],
+      "batch_size": 1,
+      "_comment": "that's all"
+    },
+    "numb_steps": 200,
+    "gradient_max_norm": 5.0,
+    "seed": 10,
+    "disp_file": "lcurve.out",
+    "disp_freq": 100,
+    "save_freq": 2000,
+    "_comment": "that's all"
+  }
+}
diff --git a/examples/density/dpa3/input.json b/examples/density/dpa3/input.json
new file mode 100644
index 0000000000..f23f5c6978
--- /dev/null
+++ b/examples/density/dpa3/input.json
@@ -0,0 +1,108 @@
+{
+  "_comment": "that's all",
+  "model": {
+    "type_map": [
+      "Li",
+      "Ni",
+      "Co",
+      "Mn",
+      "O",
+      "C",
+      "H",
+      "N",
+      "F",
+      "X"
+    ],
+    "descriptor": {
+      "type": "dpa3",
+      "repflow": {
+        "n_dim": 128,
+        "e_dim": 64,
+        "a_dim": 32,
+        "nlayers": 12,
+        "e_rcut": 6.0,
+        "e_rcut_smth": 5.3,
+        "e_sel": 120,
+        "a_rcut": 4.0,
+        "a_rcut_smth": 3.5,
+        "a_sel": 30,
+        "axis_neuron": 4,
+        "skip_stat": true,
+        "a_compress_rate": 1,
+        "a_compress_e_rate": 2,
+        "a_compress_use_split": true,
+        "update_angle": true,
+        "smooth_edge_update": true,
+        "use_dynamic_sel": true,
+        "sel_reduce_factor": 10.0,
+        "edge_use_dist": true,
+        "use_env_envelope": true,
+        "optim_update": true,
+        "update_style": "res_residual",
+        "update_residual": 0.1,
+        "update_residual_init": "const"
+      },
+      "activation_function": "custom_silu:10.0",
+      "use_tebd_bias": false,
+      "env_protection": 0.1,
+      "precision": "float32",
+      "concat_output_tebd": false
+    },
+    "fitting_net": {
+      "type": "density",
+      "neuron": [
+        240,
+        240,
+        240
+      ],
+      "resnet_dt": true,
+      "seed": 1,
+      "precision": "float32",
+      "activation_function": "silu",
+      "_comment": " that's all"
+    },
+    "_comment": " that's all"
+  },
+  "learning_rate": {
+    "type": "exp",
+    "decay_steps": 5000,
+    "start_lr": 0.001,
+    "stop_lr": 1e-05,
+    "_comment": "that's all"
+  },
+  "loss": {
+    "type": "grid_density",
+    "start_pref_d": 1,
+    "limit_pref_d": 1,
+    "_comment": " that's all"
+  },
+  "training": {
+    "stat_file": "./qm9_charge_density.hdf5",
+    "training_data": {
+      "systems": [
+        "../dataset/qm9/C7H15NO_train"
+      ],
+      "batch_size": "auto:128",
+      "_comment": "that's all"
+    },
+    "validation_data": {
+      "systems": [
+        "../dataset/qm9/C7H15NO_val"
+      ],
+      "batch_size": 1,
+      "numb_btch": 3,
+      "_comment": "that's all"
+    },
+    "numb_steps": 100,
+    "gradient_max_norm": 5.0,
+    "seed": 10,
+    "max_ckpt_keep": 100,
+    "disp_file": "lcurve.out",
+    "disp_freq": 100,
+    "save_freq": 10000,
+    "_comment": "that's all"
+  },
+  "optimizer": {
+    "type": "AdamW"
+  }
+}
diff --git a/examples/density/dptest_density_script.py b/examples/density/dptest_density_script.py
new file mode 100644
index 0000000000..9dfc2d7351
--- /dev/null
+++ b/examples/density/dptest_density_script.py
@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Evaluate a charge-density model on validation/test datasets.
+
+Usage
+-----
+python dptest_density_script.py model.pt /path/to/data --ratio 0.1 --output result.txt
+"""
+
+import argparse
+import glob
+import math
+import os
+import random
+import sys
+
+import numpy as np
+
+from deepmd.infer import (
+    DeepEval,
+)
+from deepmd.utils.data import (
+    DeepmdData,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Evaluate a charge-density model on deepmd/npy datasets."
+    )
+    parser.add_argument("model", type=str, help="Path to the model file (.pt or .pth).")
+    parser.add_argument(
+        "data_dir", type=str, help="Root directory of deepmd/npy datasets."
+    )
+    parser.add_argument(
+        "--ratio",
+        type=float,
+        default=0.1,
+        help="Fraction of frames to randomly sample from each system (default: 0.1).",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default=None,
+        help="If provided, write screen output to this file as well.",
+    )
+    parser.add_argument(
+        "--pred-file",
+        type=str,
+        default="result.d.out",
+        help="File to save paired [prediction, label] array (default: result.d.out).",
+    )
+    return parser.parse_args()
+
+
+class TeeLogger:
+    """Redirect stdout to both the terminal and a file."""
+
+    def __init__(self, filepath: str) -> None:
+        self.terminal = sys.stdout
+        self.log = open(filepath, "w")
+
+    def write(self, message: str) -> None:
+        self.terminal.write(message)
+        self.log.write(message)
+
+    def flush(self) -> None:
+        self.terminal.flush()
+        self.log.flush()
+
+    def close(self) -> None:
+        self.log.close()
+
+
+def evaluate(
+    model_path: str, data_dir: str, ratio: float
+) -> tuple[np.ndarray, np.ndarray]:
+    """Run inference and return (predictions, labels)."""
+    if not 0.0 < ratio <= 1.0:
+        raise ValueError(f"ratio must be in (0, 1], got {ratio}")
+    # DeepEval dispatches density models to DeepDensity automatically.
+    dm = DeepEval(model_path)
+    type_map = dm.get_type_map()
+
+    pred_list, label_list = [], []
+    pattern = os.path.join(data_dir, "**/type.raw")
+    systems = glob.glob(pattern, recursive=True)
+
+    if not systems:
+        raise RuntimeError(f"No deepmd/npy systems found under '{data_dir}'")
+
+    for f in systems:
+        sys_name = os.path.dirname(f)
+        data = DeepmdData(sys_name, type_map=type_map)
+        data.add(
+            "grid",
+            3,
+            atomic=False,
+            must=True,
+            high_prec=True,
+            special_shape="frame_major",
+        )
+        data.add(
+            "density",
+            1,
+            atomic=False,
+            must=True,
+            high_prec=True,
+            special_shape="frame_major",
+        )
+        test_data = data.get_test()
+
+        n_frames = test_data["coord"].shape[0]
+        n_sample = max(1, math.floor(n_frames * ratio))
+        indices = random.sample(range(n_frames), n_sample)
+
+        coord = test_data["coord"][indices].reshape(n_sample, -1)
+        if data.mixed_type:
+            atype = test_data["type"][indices].reshape(n_sample, -1)
+        else:
+            atype = test_data["type"][0]
+        cell = test_data["box"][indices] if data.pbc else None
+        grid = test_data["grid"][indices].reshape(n_sample, -1, 3)
+
+        density_pred = dm.eval(
+            coord,
+            cell,
+            atype,
+            grid=grid,
+            mixed_type=data.mixed_type,
+        )
+        density_label = test_data["density"][indices].reshape(n_sample, -1)
+
+        # Flatten each system before concatenating so systems may carry
+        # different grid extents.
+        pred_list.append(density_pred.reshape(-1))
+        label_list.append(density_label.reshape(-1))
+        sys.stdout.write(f"  {sys_name:60s}  frames={n_sample}/{n_frames}" + "\n")
+
+    predictions = np.concatenate(pred_list)
+    labels = np.concatenate(label_list)
+    return predictions, labels
+
+
+def print_summary(pred: np.ndarray, label: np.ndarray) -> None:
+    diff = pred - label
+    mae = np.mean(np.abs(diff))
+    rmse = np.sqrt(np.mean(diff**2))
+    label_mean_abs = np.mean(np.abs(label))
+    label_std = np.std(label)
+
+    sys.stdout.write("\n" + "=" * 60 + "\n")
+    sys.stdout.write("Summary" + "\n")
+    sys.stdout.write("=" * 60 + "\n")
+    sys.stdout.write(f"  Number of grid points : {label.size}" + "\n")
+    sys.stdout.write(f"  Label std             : {label_std:.6e}" + "\n")
+    sys.stdout.write(f"  RMSE                  : {rmse:.6e}" + "\n")
+    sys.stdout.write(f"  MAE                   : {mae:.6e}" + "\n")
+    sys.stdout.write(f"  Mean |label|          : {label_mean_abs:.6e}" + "\n")
+    sys.stdout.write(
+        f"  epsilon_MAE (MAE/Mean|label|) : {mae / label_mean_abs:.6e}" + "\n"
+    )
+    sys.stdout.write("=" * 60 + "\n")
+
+
+def main() -> None:
+    args = parse_args()
+
+    if args.output:
+        tee = TeeLogger(args.output)
+        sys.stdout = tee
+        sys.stdout.write(f"[INFO] Screen output will also be saved to: {args.output}\n")
+
+    sys.stdout.write(f"[INFO] Model : {args.model}" + "\n")
+    sys.stdout.write(f"[INFO] Data  : {args.data_dir}" + "\n")
+    sys.stdout.write(f"[INFO] Ratio : {args.ratio}\n")
+
+    pred, label = evaluate(args.model, args.data_dir, args.ratio)
+
+    # Save paired predictions & labels
+    out_array = np.stack([pred.reshape(-1), label.reshape(-1)], axis=1)
+    np.savetxt(args.pred_file, out_array)
+    sys.stdout.write(f"\n[INFO] Paired [pred, label] saved to: {args.pred_file}" + "\n")
+
+    print_summary(pred, label)
+
+    if args.output:
+        tee.close()
+        sys.stdout = tee.terminal
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/spin/dpa4/input-deepspin.json b/examples/spin/dpa4/input-deepspin.json
deleted file mode 100644
index c3824c947e..0000000000
--- a/examples/spin/dpa4/input-deepspin.json
+++ /dev/null
@@ -1,119 +0,0 @@
-{
-  "_comment": "DPA4/SeZM spin-energy training example using the virtual-atom (DeepSpin) scheme; see input.json for the native scheme.",
-  "model": {
-    "type": "DPA4",
-    "type_map": [
-      "Ni",
-      "O"
-    ],
-    "spin": {
-      "scheme": "deepspin",
-      "virtual_scale": [
-        0.314
-      ],
-      "use_spin": [
-        true,
-        false
-      ]
-    },
-    "descriptor": {
-      "sel": 120,
-      "rcut": 6.0,
-      "channels": 32,
-      "n_radial": 16,
-      "use_env_seed": true,
-      "lmax": 3,
-      "mmax": 1,
-      "n_blocks": 2,
-      "so2_layers": 3,
-      "radial_so2_mode": "degree_channel",
-      "radial_so2_rank": 1,
-      "n_focus": 2,
-      "focus_dim": 0,
-      "n_atten_head": 1,
-      "ffn_neurons": 0,
-      "ffn_so3_grid": true,
-      "grid_mlp": false,
-      "grid_branch": 1,
-      "ffn_blocks": 2,
-      "sandwich_norm": [
-        false,
-        true,
-        true,
-        false
-      ],
-      "message_node_so3": true,
-      "use_amp": true,
-      "precision": "float32",
-      "seed": 42
-    },
-    "fitting_net": {
-      "neuron": [
-        0
-      ],
-      "precision": "float32",
-      "seed": 42
-    },
-    "use_compile": false,
-    "enable_tf32": true
-  },
-  "learning_rate": {
-    "type": "wsd",
-    "start_lr": 4.5e-4,
-    "stop_lr": 1e-6,
-    "warmup_steps": 5000,
-    "warmup_start_factor": 0.2,
-    "decay_phase_ratio": 0.65,
-    "decay_type": "cosine"
-  },
-  "loss": {
-    "type": "ener_spin",
-    "start_pref_e": 0.02,
-    "limit_pref_e": 1,
-    "start_pref_fr": 1000,
-    "limit_pref_fr": 1,
-    "start_pref_fm": 1000,
-    "limit_pref_fm": 1
-  },
-  "optimizer": {
-    "type": "HybridMuon",
-    "weight_decay": 0.001
-  },
-  "training": {
-    "stat_file": "./dpa4_spin.hdf5",
-    "training_data": {
-      "systems": [
-        "../data_reformat/data_0",
-        "../data_reformat/data_1"
-      ],
-      "batch_size": 1
-    },
-    "validation_data": {
-      "systems": [
-        "../data_reformat/data_2"
-      ],
-      "batch_size": 1,
-      "numb_batch": 1
-    },
-    "numb_steps": 2000000,
-    "gradient_max_norm": 5.0,
-    "save_freq": 2000,
-    "max_ckpt_keep": 3,
-    "enable_ema": true,
-    "ema_decay": 0.999,
-    "ema_ckpt_keep": 3,
-    "disp_file": "lcurve.out",
-    "disp_freq": 1000,
-    "disp_avg": true,
-    "disp_training": true,
-    "time_training": true,
-    "tensorboard": false,
-    "enable_profiler": false,
-    "tensorboard_freq": 1000,
-    "tensorboard_log_dir": "tb_log",
-    "profiling": false,
-    "profiling_file": "timeline.json",
-    "zero_stage": 1,
-    "seed": 42
-  }
-}
diff --git a/examples/spin/dpa4/lmp/README.md b/examples/spin/dpa4/lmp/README.md
index ac1a75712c..802c7ff550 100644
--- a/examples/spin/dpa4/lmp/README.md
+++ b/examples/spin/dpa4/lmp/README.md
@@ -1,8 +1,7 @@
 # LAMMPS example for DPA4 / SeZM native spin
 
 Runs a native-spin DPA4 / SeZM model in LAMMPS through `pair_style deepspin`
-backed by an AOTInductor `.pt2` archive. For the classical DeepSpin
-(virtual-atom) scheme, see `examples/spin/lmp`.
+backed by an AOTInductor `.pt2` archive.
 
 ## Files
 
diff --git a/examples/spin/dpa4c/lmp/README.md b/examples/spin/dpa4c/lmp/README.md
index aa4153f4f6..5706fcf193 100644
--- a/examples/spin/dpa4c/lmp/README.md
+++ b/examples/spin/dpa4c/lmp/README.md
@@ -24,9 +24,9 @@ dp --pt-expt train ../input.json
 dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph
 ```
 
-The pair style accepts either a graph-lower archive, as frozen above, or the
-compact canonical graph archive that `dp --pt-expt compress` produces, which
-runs the fused CUDA kernels. See
+`dpa4spin` accepts either a graph-lower archive, as frozen above, or the
+compressed archive that `dp --pt-expt compress` produces and that runs the
+fused CUDA kernels. See
 [model compression](../../../../doc/model/dpa4c.md#model-compression) for that
 workflow, and point `pair_style` at the compressed archive if you follow it.
 
@@ -36,8 +36,9 @@ Run on the host:
 lmp -in in.lammps
 ```
 
-Run device-resident under Kokkos, where the graph and the moment stay in device
-memory for the whole step:
+Under Kokkos the graph and the moments stay in device memory for the whole
+step, and `dpa4spin/kk` reads only the compressed archive, so `pair_style` has
+to point at that archive. Run device-resident with:
 
 ```bash
 lmp -k on g 1 -sf kk -in in.lammps
@@ -64,6 +65,17 @@ configuration, which exercises the conservative force, the neighbor rebuilds
 and the domain decomposition; the thermostat does not touch the moments, so
 they stay fixed through that stage.
 
-Spin dynamics through `fix nve/spin` requires a LAMMPS build whose fix
-recognizes this pair style, because the stock fix accumulates the magnetic
-force only from pair styles matching its own name pattern.
+By default the pair style stores the precession magnetic force that the LAMMPS
+SPIN package expects, the negative energy gradient `-dE/dS` scaled by the
+moment magnitude over `hbar`. The `fix nve/spin` of a stock LAMMPS build does
+not use this force: the fix evaluates the magnetic force itself, through the
+pair styles whose names begin with `spin`, and `dpa4spin` and `dpa4spin/kk` are
+not among them. With no such pair style, the fix runs silently on a single MPI
+rank and stops with a sectoring error on several. Spin dynamics therefore needs
+a LAMMPS build whose integrator reads the magnetic force of the pair style.
+
+A Newtonian integrator of the Cartesian spin vector needs `-dE/dS` itself
+rather than the precession force. Appending `magnetic_force cartesian` to the
+`pair_style` line selects `-dE/dS` and requires `units metal`; the pair style
+declares the selected convention through
+`Pair::extract("fm_cartesian_force", dim)`.
diff --git a/examples/spin/dpa4c/lmp/in.lammps b/examples/spin/dpa4c/lmp/in.lammps
index 050ce53903..bcbe235ba4 100644
--- a/examples/spin/dpa4c/lmp/in.lammps
+++ b/examples/spin/dpa4c/lmp/in.lammps
@@ -4,13 +4,21 @@
 # magnetic force is the negative spin gradient of the energy, so no virtual
 # atoms are created and the atom count equals the number of physical atoms.
 #
-# The pair style reads a graph-lower archive or the compact canonical graph
-# archive produced by compression. See README.md for how to obtain one.
-#
-# Substitute `dpa4spin/kk` for the device-resident Kokkos path, which keeps the
-# graph and the moment in device memory:
+# `dpa4spin` reads a graph-lower archive or the compressed archive; see
+# README.md for how to obtain either. The `-sf kk` option selects the Kokkos
+# variant `dpa4spin/kk`, which keeps the graph and the moments in device memory
+# and reads only the compressed archive:
 #
 #   lmp -k on g 1 -sf kk -in in.lammps
+#
+# By default atom->fm holds the precession magnetic force. A Newtonian
+# integrator of the Cartesian spin vector needs the negative energy gradient
+# -dE/dS instead, which `magnetic_force cartesian` selects under units metal:
+#
+#   pair_style        dpa4spin/kk compressed_model.pt2 magnetic_force cartesian
+#
+# The pair style declares the selected convention through
+# Pair::extract("fm_cartesian_force", dim).
 
 units             metal
 dimension         3
@@ -42,10 +50,13 @@ dump_modify       1 sort id format float %.12e
 # magnetic force of the antiferromagnetic reference state.
 run               0
 
-# `min_style spin` reads the magnetic force from the pair style directly and
-# relaxes the moment directions at fixed positions. Note that `fix nve/spin`
-# accumulates the magnetic force only from pair styles whose name it matches,
-# so spin dynamics needs a LAMMPS build aware of this style.
+# `min_style spin` reads the magnetic force from the pair style and relaxes the
+# moment directions at fixed positions. The `fix nve/spin` of a stock LAMMPS
+# build does not use this force: the fix evaluates the magnetic force through
+# the pair styles whose names begin with `spin`, so with this pair style alone
+# the fix runs silently on a single rank and stops with a sectoring error on
+# several. Spin dynamics needs a LAMMPS build whose integrator reads the
+# magnetic force of the pair style.
 min_style         spin
 min_modify        line spin_cubic discrete_factor 10.0
 minimize          1.0e-10 1.0e-8 500 1000
diff --git a/pyproject.toml b/pyproject.toml
index 5da7f9b1c6..774cd08fa5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -79,6 +79,7 @@ repository = "https://github.com/deepmodeling/deepmd-kit"
 # which can be read by the build backend.
 [tool.deepmd_build_backend.optional-dependencies]
 test = [
+    "dependency_groups",
     "dpdata>=0.2.7",
     # ASE issue: https://gitlab.com/ase/ase/-/merge_requests/2843
     # fixed in 3.23.0
diff --git a/skills/deepmd-install/SKILL.md b/skills/deepmd-install/SKILL.md
index 89245bf0bc..6826b8f15c 100644
--- a/skills/deepmd-install/SKILL.md
+++ b/skills/deepmd-install/SKILL.md
@@ -216,6 +216,9 @@ LAMMPS, use a packaged `lmp` when it satisfies the request. Otherwise follow
 the matching built-in or plugin instructions after the C/C++ interface
 succeeds. Enable Kokkos only when the requested LAMMPS runtime requires it,
 and select the architecture supported by that exact LAMMPS/Kokkos source tree.
+For LAMMPS runs on several GPUs, build LAMMPS and the DeePMD-kit C/C++
+interface against the same CUDA-aware MPI library, such as Open MPI with UCX,
+and enable its CUDA support at run time as the matching LAMMPS page describes.
 
 ## 4. Verify the requested interface
 
@@ -233,8 +236,10 @@ An installation is complete only after the requested public interface runs:
 1. For C/C++, confirm the installed headers/libraries, then load a built
    library or run a linked client so the platform loader resolves its dynamic
    dependencies.
-1. For LAMMPS, run the selected binary with `-h`, require the exact DeePMD pair
-   style needed by the model, and run a short documented example.
+1. For LAMMPS, run the selected binary with `-h`; in plugin mode, also list the
+   loaded plugins with the `plugin list` command, because `-h` shows only the
+   built-in styles. Require the exact DeePMD pair style needed by the model, and
+   run a short documented example.
 1. For Docker, perform all applicable checks inside the selected image.
 
 On failure, stop at the first failing check and read
diff --git a/source/api_c/CMakeLists.txt b/source/api_c/CMakeLists.txt
index b11eee41a9..452f7b73c7 100644
--- a/source/api_c/CMakeLists.txt
+++ b/source/api_c/CMakeLists.txt
@@ -1,6 +1,7 @@
 file(GLOB LIB_SRC src/*.cc src/*.cpp)
 # only install public headers
 file(GLOB INC_SRC include/c_api.h include/deepmd.hpp)
+list(APPEND INC_SRC ../api_cc/include/canonical_graph.h)
 
 set(libname "${LIB_DEEPMD_C}")
 
@@ -9,8 +10,10 @@ add_library(${libname} SHARED ${LIB_SRC})
 # link: libdeepmd libdeepmd_op libtensorflow_cc libtensorflow_framework
 target_link_libraries(${libname} PRIVATE ${LIB_DEEPMD_CC})
 target_include_directories(
-  ${libname} PUBLIC $
-                    $)
+  ${libname}
+  PUBLIC $
+         $
+         $)
 
 set_target_properties(${libname} PROPERTIES INSTALL_RPATH
                                             "$ORIGIN;${BACKEND_LIBRARY_PATH}")
diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h
index d98d744387..ff8903d341 100644
--- a/source/api_c/include/c_api.h
+++ b/source/api_c/include/c_api.h
@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 #pragma once
 #include 
+
+#include "canonical_graph.h"
 #ifdef __cplusplus
 extern "C" {
 #else
@@ -13,7 +15,7 @@ extern "C" {
 /** C API version. Bumped whenever the API is changed.
  * @since API version 22
  */
-#define DP_C_API_VERSION 30
+#define DP_C_API_VERSION 31
 
 /**
  * @brief Neighbor list.
@@ -427,39 +429,101 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp,
                                              const DP_Nlist* comm_nlist);
 
 /**
- * @brief Evaluate a compact canonical graph on the model device.
+ * @brief Evaluate a compact canonical graph on the model device: the
+ * device-resident inference of the DPA4 family, whose compressed DPA4C
+ * archives are frozen with the compact canonical lower, as are compressed
+ * DPA1 archives.
  *
  * @param[in] dp The DP to use.
  * @param[out] d_atom_energy Per-local-atom energy, shape ``(nloc)``.
  * @param[out] d_force Per-node force, shape ``(nall_nodes, 3)``.
- * @param[out] d_atom_virial Per-node virial, shape ``(nall_nodes, 9)``.
+ * @param[out] d_virial Virial of the graph, nine components, or NULL when it
+ * is not requested.
+ * @param[out] d_atom_virial Per-node virial, shape ``(nall_nodes, 9)``, or
+ * NULL when the per-atom virial is not requested.
  * @param[in] d_atype Per-node atom types, shape ``(nall_nodes)``.
  * @param[in] d_source Source node per edge storage slot.
+ * @param[in] index_width Width of each source index, four or eight bytes.
  * @param[in] d_edge_vec FP32 edge vectors, shape ``(edge_storage, 3)``.
  * @param[in] d_destination_row_ptr Destination CSR offsets.
- * @param[in] d_source_row_ptr Source CSR offsets.
- * @param[in] d_source_order Source-grouped edge storage positions.
+ * @param[in] d_source_row_ptr Source CSR offsets. May be NULL together with
+ * ``d_source_order`` and ``d_atom_virial`` unless
+ * DP_DeepPotCanonicalGraphRequiresSourceCSR reports otherwise.
+ * @param[in] d_source_order Source-grouped int64 edge storage positions.
+ * @param[in] d_coord Node positions in the length unit of the model, shape
+ * ``(nall_nodes, 3)``, in the order of ``d_atype``, or NULL. Only a model
+ * with a long-range module (DP_DeepPotHasLongRange) reads them, and it
+ * requires them; every other model ignores them.
+ * @param[in] world Pointer to the MPI communicator of the caller (an
+ * ``MPI_Comm``), as the ``world`` argument of DP_NewNlist_comm, when the
+ * system spans several ranks, or NULL on a single rank. A model with a
+ * long-range module evaluates a complete system: a single rank holds every
+ * node of a non-periodic system (``nall_nodes == nloc``), and the halo nodes
+ * of a periodic one are images of its owned nodes; on several ranks it
+ * requires the communicator, reads the positions of the nodes each rank owns,
+ * and is evaluated on every rank at every step, a rank without nodes included
+ * (``nall_nodes``, ``nloc`` and ``edge_storage`` zero), whose buffers sized by
+ * the nodes or by the edge storage may then be NULL. Every other model
+ * ignores it.
+ * @param[in] box The cell of a periodic system in host memory, nine values in
+ * the length unit of the model, row-major with the lattice vectors as rows,
+ * or NULL for a non-periodic system. A model with a long-range module that
+ * serves periodic systems evaluates the system over the images of the cell,
+ * one whose module serves none rejects a cell, and every other model ignores
+ * it.
  * @param[in] nloc Number of owned local nodes.
  * @param[in] nall_nodes Total local-plus-halo node count.
  * @param[in] edge_storage Number of edge storage slots.
- * @note The source and source-order arrays are uint32 as of API version 30;
- *       API versions 28 and 29 declared them as signed int64.
- * @since API version 28
- */
-extern void DP_DeepPotComputeCanonicalGraphGPU(
+ * @param[in] stream Accelerator stream (``cudaStream_t``) of the model device
+ * on which the inputs were produced and the outputs will be consumed; a
+ * non-NULL stream requires a CUDA build. The evaluation is enqueued on it
+ * without a synchronization of the stream. A graph without halo nodes is
+ * captured into a CUDA graph on the second consecutive call with the same
+ * buffers and extents and replayed while they stay the same, except where a
+ * model with a long-range module receives a communicator. A model whose module
+ * serves periodic systems copies the cell into memory of its own, which a
+ * replay reads, and the mesh size of the cell counts as an extent: a call with
+ * another cell than the previous call first waits for the earlier evaluations,
+ * consecutive evaluations on different streams run in call order, and after the
+ * mesh size changes, the model evaluates without replay for 64 calls, after
+ * which the capture follows as for new extents. The host waits for the device
+ * when a call captures or releases a graph, which happens on that second call
+ * and on a call whose buffers or extents change. DP_DeepPotSetChargeSpin waits
+ * for the work in flight on the model device before it replaces the constants
+ * and for the replacement before it returns, and the caller completes the work
+ * on its stream before it deletes the model. NULL runs on the current stream of
+ * the calling thread for the model device and waits for that stream before
+ * returning. Models in different threads of one process evaluate concurrently,
+ * and their captures, releases of graphs and charge-state updates, which wait
+ * for the whole device, take turns; CUDA work outside DeePMD-kit that waits for
+ * the whole device while a capture is under way invalidates it, and the model
+ * then evaluates without replay.
+ * @note API versions 28 to 30 named this function
+ *       DP_DeepPotComputeCanonicalGraphGPU, without ``d_virial``, ``d_coord``,
+ *       ``world``, ``box`` and ``stream``; a binary built against that name
+ *       fails to load rather than call a function of another argument list.
+ * @since API version 31
+ */
+extern void DP_DPA4ComputeCanonicalGraphGPU(
     DP_DeepPot* dp,
     double* d_atom_energy,
     double* d_force,
+    double* d_virial,
     double* d_atom_virial,
     const int64_t* d_atype,
-    const uint32_t* d_source,
+    const void* d_source,
+    DP_CanonicalIndexWidth index_width,
     const float* d_edge_vec,
     const int64_t* d_destination_row_ptr,
     const int64_t* d_source_row_ptr,
-    const uint32_t* d_source_order,
-    int nloc,
-    int nall_nodes,
-    int64_t edge_storage);
+    const int64_t* d_source_order,
+    const double* d_coord,
+    void* world,
+    const double* box,
+    int64_t nloc,
+    int64_t nall_nodes,
+    int64_t edge_storage,
+    void* stream);
 
 /**
  * @brief Query whether the loaded artifact supports device-edge inference.
@@ -479,6 +543,25 @@ extern bool DP_DeepPotUsesFP32EdgeVectors(DP_DeepPot* dp);
  */
 extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp);
 
+/**
+ * @brief Query whether every compact canonical evaluation reads the source
+ * CSR, rather than only an evaluation that requests the per-atom virial.
+ * @since API version 31
+ */
+extern bool DP_DeepPotCanonicalGraphRequiresSourceCSR(DP_DeepPot* dp);
+
+/**
+ * @brief Query whether the model carries a long-range module.
+ *
+ * Such a module sums its interaction over every pair of atoms, and over the
+ * images of the cell of a periodic system when it serves periodic systems, so
+ * the model evaluates complete systems, and the device-resident entry
+ * DP_DPA4ComputeCanonicalGraphGPU passes it the node positions, the cell and,
+ * on several ranks, the communicator.
+ * @since API version 31
+ */
+extern bool DP_DeepPotHasLongRange(DP_DeepPot* dp);
+
 /**
  * @brief Evaluate the energy, force and virial by using a DP. (double version)
  * @version 2
@@ -1259,42 +1342,59 @@ extern void DP_DeepSpinComputeNListf3(DP_DeepSpin* dp,
 
 /**
  * @brief Evaluate a compact canonical graph on the model device with a DP spin
- * model.
+ * model: the device-resident inference of native-spin DPA4C archives.
  *
  * @param[in] dp The DP spin model to use.
  * @param[out] d_atom_energy Per-local-atom energy, shape ``(nloc)``.
  * @param[out] d_force Per-node force, shape ``(nall_nodes, 3)``.
  * @param[out] d_force_mag Per-node magnetic force, shape ``(nall_nodes, 3)``.
- * @param[out] d_atom_virial Per-node virial, shape ``(nall_nodes, 9)``.
+ * @param[out] d_virial Virial of the graph, nine components, or NULL when it
+ * is not requested.
+ * @param[out] d_atom_virial Per-node virial, shape ``(nall_nodes, 9)``, or
+ * NULL when the per-atom virial is not requested.
  * @param[in] d_atype Per-node atom types, shape ``(nall_nodes)``.
  * @param[in] d_source Source node per edge storage slot.
+ * @param[in] index_width Width of each source index, four or eight bytes.
  * @param[in] d_edge_vec FP32 edge vectors, shape ``(edge_storage, 3)``.
  * @param[in] d_destination_row_ptr Destination CSR offsets.
- * @param[in] d_source_row_ptr Source CSR offsets.
- * @param[in] d_source_order Source-grouped edge storage positions.
+ * @param[in] d_source_row_ptr Source CSR offsets, read only for the per-atom
+ * virial. May be NULL together with ``d_source_order`` when ``d_atom_virial``
+ * is NULL.
+ * @param[in] d_source_order Source-grouped int64 edge storage positions.
  * @param[in] d_spin FP32 per-node magnetic moment, shape ``(nall_nodes, 3)``;
  * halo rows carry their owner's moment.
  * @param[in] nloc Number of owned local nodes.
  * @param[in] nall_nodes Total local-plus-halo node count.
  * @param[in] edge_storage Number of edge storage slots.
- * @since API version 30
- */
-extern void DP_DeepSpinComputeCanonicalGraphGPU(
+ * @param[in] stream Accelerator stream (``cudaStream_t``) of the model device
+ * on which the inputs were produced and the outputs will be consumed, under
+ * the contract described for DP_DPA4ComputeCanonicalGraphGPU.
+ * DP_DeepSpinSetChargeSpin waits for the work in flight on the model device
+ * before it replaces the constants and for the replacement before it returns.
+ * @note API version 30 named this function DP_DeepSpinComputeCanonicalGraphGPU,
+ *       without ``d_virial`` and ``stream``; a binary built against that name
+ *       fails to load rather than call a function of another argument list.
+ * @since API version 31
+ */
+extern void DP_DPA4SpinComputeCanonicalGraphGPU(
     DP_DeepSpin* dp,
     double* d_atom_energy,
     double* d_force,
     double* d_force_mag,
+    double* d_virial,
     double* d_atom_virial,
     const int64_t* d_atype,
-    const uint32_t* d_source,
+    const void* d_source,
+    DP_CanonicalIndexWidth index_width,
     const float* d_edge_vec,
     const int64_t* d_destination_row_ptr,
     const int64_t* d_source_row_ptr,
-    const uint32_t* d_source_order,
+    const int64_t* d_source_order,
     const float* d_spin,
-    int nloc,
-    int nall_nodes,
-    int64_t edge_storage);
+    int64_t nloc,
+    int64_t nall_nodes,
+    int64_t edge_storage,
+    void* stream);
 
 /**
  * @brief Query whether the compact canonical graph ABI is active for a DP spin
@@ -2683,6 +2783,14 @@ int DP_DeepPotModelDeviGetDimAParam(DP_DeepPotModelDevi* dp);
  */
 int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp);
 
+/**
+ * @brief Query whether any model of a DP Model Deviation carries a long-range
+ * module; see DP_DeepPotHasLongRange.
+ * @param[in] dp The DP Model Deviation to use.
+ * @since API version 31
+ */
+extern bool DP_DeepPotModelDeviHasLongRange(DP_DeepPotModelDevi* dp);
+
 /**
  * @brief Fix the charge/spin condition served for the rest of the run.
  * Applied to every model, so that the deviation is taken between models under
diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp
index 681b197b0e..7fe1560b05 100644
--- a/source/api_c/include/deepmd.hpp
+++ b/source/api_c/include/deepmd.hpp
@@ -1359,24 +1359,61 @@ class DeepPot : public DeepBaseModel {
 
   /**
    * @brief Evaluate a compact canonical graph on the model device.
+   *
+   * ``d_virial`` may be null when the virial is not requested.
+   * ``d_atom_virial`` may be null, and then so may the source CSR unless
+   * canonical_graph_requires_source_csr() holds. ``d_coord`` holds the node
+   * positions, ``world`` the MPI communicator of a system on several ranks
+   * and ``box`` the cell of a periodic system, which a model with a
+   * long-range module reads (has_long_range()); all three may be null for
+   * any other model. A non-null ``stream`` is the
+   * caller's accelerator stream on the model device. The arguments have the
+   * semantics of DP_DPA4ComputeCanonicalGraphGPU.
    */
   void compute_canonical_graph_gpu(double* d_atom_energy,
                                    double* d_force,
+                                   double* d_virial,
                                    double* d_atom_virial,
                                    const int64_t* d_atype,
-                                   const uint32_t* d_source,
+                                   const void* d_source,
+                                   DP_CanonicalIndexWidth index_width,
                                    const float* d_edge_vec,
                                    const int64_t* d_destination_row_ptr,
                                    const int64_t* d_source_row_ptr,
-                                   const uint32_t* d_source_order,
-                                   const int nloc,
-                                   const int nall_nodes,
-                                   const int64_t edge_storage) {
-    DP_DeepPotComputeCanonicalGraphGPU(
-        dp, d_atom_energy, d_force, d_atom_virial, d_atype, d_source,
-        d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order,
-        nloc, nall_nodes, edge_storage);
+                                   const int64_t* d_source_order,
+                                   const double* d_coord,
+                                   void* world,
+                                   const double* box,
+                                   const int64_t nloc,
+                                   const int64_t nall_nodes,
+                                   const int64_t edge_storage,
+                                   void* stream = nullptr) {
+    DP_DPA4ComputeCanonicalGraphGPU(
+        dp, d_atom_energy, d_force, d_virial, d_atom_virial, d_atype, d_source,
+        index_width, d_edge_vec, d_destination_row_ptr, d_source_row_ptr,
+        d_source_order, d_coord, world, box, nloc, nall_nodes, edge_storage,
+        stream);
+    DP_CHECK_OK(DP_DeepPotCheckOK, dp);
+  }
+
+  /**
+   * @brief Query whether every compact canonical evaluation reads the source
+   * CSR.
+   */
+  bool canonical_graph_requires_source_csr() const {
+    const bool result = DP_DeepPotCanonicalGraphRequiresSourceCSR(dp);
+    DP_CHECK_OK(DP_DeepPotCheckOK, dp);
+    return result;
+  }
+
+  /**
+   * @brief Query whether the model carries a long-range module, which
+   * evaluates complete systems; see DP_DeepPotHasLongRange.
+   */
+  bool has_long_range() const {
+    const bool result = DP_DeepPotHasLongRange(dp);
     DP_CHECK_OK(DP_DeepPotCheckOK, dp);
+    return result;
   }
 
   /**
@@ -1919,26 +1956,34 @@ class DeepSpin : public DeepBaseModel {
    * @brief Evaluate a compact canonical graph on the model device.
    *
    * Graph, moment, and output pointers reside on the model device; the halo
-   * rows of the moment carry their owner's value.
+   * rows of the moment carry their owner's value. ``d_virial`` may be null
+   * when the virial is not requested. ``d_atom_virial`` may be null, and then
+   * so may the source CSR. A non-null ``stream`` is the caller's accelerator
+   * stream on the model device, with the semantics of
+   * DP_DPA4ComputeCanonicalGraphGPU.
    */
   void compute_canonical_graph_gpu(double* d_atom_energy,
                                    double* d_force,
                                    double* d_force_mag,
+                                   double* d_virial,
                                    double* d_atom_virial,
                                    const int64_t* d_atype,
-                                   const uint32_t* d_source,
+                                   const void* d_source,
+                                   DP_CanonicalIndexWidth index_width,
                                    const float* d_edge_vec,
                                    const int64_t* d_destination_row_ptr,
                                    const int64_t* d_source_row_ptr,
-                                   const uint32_t* d_source_order,
+                                   const int64_t* d_source_order,
                                    const float* d_spin,
-                                   const int nloc,
-                                   const int nall_nodes,
-                                   const int64_t edge_storage) {
-    DP_DeepSpinComputeCanonicalGraphGPU(
-        dp, d_atom_energy, d_force, d_force_mag, d_atom_virial, d_atype,
-        d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr,
-        d_source_order, d_spin, nloc, nall_nodes, edge_storage);
+                                   const int64_t nloc,
+                                   const int64_t nall_nodes,
+                                   const int64_t edge_storage,
+                                   void* stream = nullptr) {
+    DP_DPA4SpinComputeCanonicalGraphGPU(
+        dp, d_atom_energy, d_force, d_force_mag, d_virial, d_atom_virial,
+        d_atype, d_source, index_width, d_edge_vec, d_destination_row_ptr,
+        d_source_row_ptr, d_source_order, d_spin, nloc, nall_nodes,
+        edge_storage, stream);
     DP_CHECK_OK(DP_DeepSpinCheckOK, dp);
   }
 
@@ -2608,6 +2653,17 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
     return dchgspin;
   }
 
+  /**
+   * @brief Query whether any of the models carries a long-range module; see
+   * DP_DeepPotModelDeviHasLongRange.
+   **/
+  bool has_long_range() const {
+    assert(dp);
+    const bool result = DP_DeepPotModelDeviHasLongRange(dp);
+    DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp);
+    return result;
+  }
+
   /**
    * @brief Fix the charge/spin condition served for the rest of the run.
    * Applied to every model, so that the deviation is taken between models
diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc
index c35a4a1763..eeb196fead 100644
--- a/source/api_c/src/c_api.cc
+++ b/source/api_c/src/c_api.cc
@@ -1817,24 +1817,31 @@ void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp,
                                     aparam_size, nall_nodes, comm_nlist);
 }
 
-void DP_DeepPotComputeCanonicalGraphGPU(DP_DeepPot* dp,
-                                        double* d_atom_energy,
-                                        double* d_force,
-                                        double* d_atom_virial,
-                                        const int64_t* d_atype,
-                                        const uint32_t* d_source,
-                                        const float* d_edge_vec,
-                                        const int64_t* d_destination_row_ptr,
-                                        const int64_t* d_source_row_ptr,
-                                        const uint32_t* d_source_order,
-                                        const int nloc,
-                                        const int nall_nodes,
-                                        const int64_t edge_storage) {
-  DP_REQUIRES_OK(dp,
-                 dp->dp.compute_canonical_graph_gpu(
-                     d_atom_energy, d_force, d_atom_virial, d_atype, d_source,
-                     d_edge_vec, d_destination_row_ptr, d_source_row_ptr,
-                     d_source_order, nloc, nall_nodes, edge_storage));
+void DP_DPA4ComputeCanonicalGraphGPU(DP_DeepPot* dp,
+                                     double* d_atom_energy,
+                                     double* d_force,
+                                     double* d_virial,
+                                     double* d_atom_virial,
+                                     const int64_t* d_atype,
+                                     const void* d_source,
+                                     DP_CanonicalIndexWidth index_width,
+                                     const float* d_edge_vec,
+                                     const int64_t* d_destination_row_ptr,
+                                     const int64_t* d_source_row_ptr,
+                                     const int64_t* d_source_order,
+                                     const double* d_coord,
+                                     void* world,
+                                     const double* box,
+                                     const int64_t nloc,
+                                     const int64_t nall_nodes,
+                                     const int64_t edge_storage,
+                                     void* stream) {
+  DP_REQUIRES_OK(
+      dp, dp->dp.compute_canonical_graph_gpu(
+              d_atom_energy, d_force, d_virial, d_atom_virial, d_atype,
+              d_source, index_width, d_edge_vec, d_destination_row_ptr,
+              d_source_row_ptr, d_source_order, d_coord, world, box, nloc,
+              nall_nodes, edge_storage, stream));
 }
 
 bool DP_DeepPotSupportsDeviceEdgeInference(DP_DeepPot* dp) {
@@ -1864,6 +1871,24 @@ bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp) {
   }
 }
 
+bool DP_DeepPotCanonicalGraphRequiresSourceCSR(DP_DeepPot* dp) {
+  try {
+    return dp->dp.canonical_graph_requires_source_csr();
+  } catch (deepmd::deepmd_exception& ex) {
+    dp->exception = std::string(ex.what());
+    return false;
+  }
+}
+
+bool DP_DeepPotHasLongRange(DP_DeepPot* dp) {
+  try {
+    return dp->dp.has_long_range();
+  } catch (deepmd::deepmd_exception& ex) {
+    dp->exception = std::string(ex.what());
+    return false;
+  }
+}
+
 // multiple frames
 void DP_DeepPotCompute2(DP_DeepPot* dp,
                         const int nframes,
@@ -2206,26 +2231,30 @@ void DP_DeepSpinComputeNListf3(DP_DeepSpin* dp,
       charge_spin);
 }
 
-void DP_DeepSpinComputeCanonicalGraphGPU(DP_DeepSpin* dp,
+void DP_DPA4SpinComputeCanonicalGraphGPU(DP_DeepSpin* dp,
                                          double* d_atom_energy,
                                          double* d_force,
                                          double* d_force_mag,
+                                         double* d_virial,
                                          double* d_atom_virial,
                                          const int64_t* d_atype,
-                                         const uint32_t* d_source,
+                                         const void* d_source,
+                                         DP_CanonicalIndexWidth index_width,
                                          const float* d_edge_vec,
                                          const int64_t* d_destination_row_ptr,
                                          const int64_t* d_source_row_ptr,
-                                         const uint32_t* d_source_order,
+                                         const int64_t* d_source_order,
                                          const float* d_spin,
-                                         const int nloc,
-                                         const int nall_nodes,
-                                         const int64_t edge_storage) {
+                                         const int64_t nloc,
+                                         const int64_t nall_nodes,
+                                         const int64_t edge_storage,
+                                         void* stream) {
   DP_REQUIRES_OK(
       dp, dp->dp.compute_canonical_graph_gpu(
-              d_atom_energy, d_force, d_force_mag, d_atom_virial, d_atype,
-              d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr,
-              d_source_order, d_spin, nloc, nall_nodes, edge_storage));
+              d_atom_energy, d_force, d_force_mag, d_virial, d_atom_virial,
+              d_atype, d_source, index_width, d_edge_vec, d_destination_row_ptr,
+              d_source_row_ptr, d_source_order, d_spin, nloc, nall_nodes,
+              edge_storage, stream));
 }
 
 bool DP_DeepSpinUsesCanonicalGraphInference(DP_DeepSpin* dp) {
@@ -2848,6 +2877,15 @@ int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp) {
   return dp->dp.dim_chg_spin();
 }
 
+bool DP_DeepPotModelDeviHasLongRange(DP_DeepPotModelDevi* dp) {
+  try {
+    return dp->dp.has_long_range();
+  } catch (deepmd::deepmd_exception& ex) {
+    dp->exception = std::string(ex.what());
+    return false;
+  }
+}
+
 void DP_DeepPotModelDeviSetChargeSpin(DP_DeepPotModelDevi* dp,
                                       const double* charge_spin,
                                       const int numb_chg_spin) {
diff --git a/source/api_c/tests/CMakeLists.txt b/source/api_c/tests/CMakeLists.txt
index f1a74a26a6..0be6b51b62 100644
--- a/source/api_c/tests/CMakeLists.txt
+++ b/source/api_c/tests/CMakeLists.txt
@@ -20,6 +20,11 @@ if(ENABLE_PYTORCH)
   # BUILD_PT_EXPT is set for the test binary; otherwise pt_expt tests all
   # GTEST_SKIP() with "PyTorch support is not enabled".
   target_link_libraries(runUnitTests_c PRIVATE "${TORCH_LIBRARIES}")
+  # The device-resident cases order their work on a caller's stream through the
+  # c10 CUDA stream API, which the backend uses under the same condition.
+  if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA)
+    target_compile_definitions(runUnitTests_c PRIVATE DEEPMD_TORCH_HAS_CUDA)
+  endif()
 endif()
 if(ENABLE_JAX)
   target_compile_definitions(runUnitTests_c PRIVATE BUILD_JAX)
diff --git a/source/api_c/tests/test_deeppot_dpa4c_canonical_hpp.cc b/source/api_c/tests/test_deeppot_dpa4c_canonical_hpp.cc
new file mode 100644
index 0000000000..c81003852e
--- /dev/null
+++ b/source/api_c/tests/test_deeppot_dpa4c_canonical_hpp.cc
@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// The device-resident entry points of the compressed DPA4C archives frozen
+// with the compact canonical lower, through deepmd::hpp::DeepPot and
+// deepmd::hpp::DeepSpin and so through DP_DPA4ComputeCanonicalGraphGPU and
+// DP_DPA4SpinComputeCanonicalGraphGPU of the C API (see
+// source/api_cc/tests/dpa4c_canonical_test_common.h for the cases).
+#include 
+
+#include "deepmd.hpp"
+#include "test_utils.h"
+// The shared cases follow the helpers of test_utils.h, which they call.
+#include "../../api_cc/tests/dpa4c_canonical_test_common.h"
+
+namespace deepmd_test {
+namespace dpa4c_canonical {
+
+using Models = ::testing::Types,
+                                Spin>;
+INSTANTIATE_TYPED_TEST_SUITE_P(C,
+                               TestInferDpa4cCanonicalPtExpt,
+                               Models,
+                               CaseNames);
+
+}  // namespace dpa4c_canonical
+}  // namespace deepmd_test
diff --git a/source/api_c/tests/test_deeppot_long_range_hpp.cc b/source/api_c/tests/test_deeppot_long_range_hpp.cc
new file mode 100644
index 0000000000..9e6804e056
--- /dev/null
+++ b/source/api_c/tests/test_deeppot_long_range_hpp.cc
@@ -0,0 +1,81 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Test the C API with a pt_expt (.pt2) DPA4C model that carries a
+// latent-charge long-range module (source/tests/infer/gen_dpa4c_long_range.py).
+#include 
+
+#include 
+#include 
+#include 
+
+#include "../../api_cc/tests/expected_ref.h"
+#include "DeepPotPTExpt.h"
+#include "deepmd.hpp"
+#include "test_utils.h"
+
+namespace {
+constexpr const char* kModelPath =
+    "../../tests/infer/deeppot_dpa4c_long_range.pt2";
+constexpr const char* kRefPath =
+    "../../tests/infer/deeppot_dpa4c_long_range.expected";
+constexpr const char* kShortRangeModelPath =
+    "../../tests/infer/deeppot_sea.pth";
+}  // namespace
+
+template 
+class TestInferDeepPotLongRangeHPP : public ::testing::Test {
+ protected:
+  std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74,
+                                  0.25,  3.32, 1.68, 3.36,  3.00, 1.81,
+                                  3.51,  2.51, 2.60, 4.27,  3.22, 1.56};
+  std::vector atype = {0, 1, 1, 0, 1, 1};
+  std::vector expected_f;
+  double expected_tot_e = 0.;
+  deepmd::hpp::DeepPot dp;
+
+  void SetUp() override {
+#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT
+    GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#endif
+    dp.init(kModelPath);
+    deepmd_test::ExpectedRef ref;
+    ref.load(kRefPath);
+    expected_f = ref.get("default", "expected_f");
+    for (const VALUETYPE value : ref.get("default", "expected_e")) {
+      expected_tot_e += value;
+    }
+  }
+};
+
+TYPED_TEST_SUITE(TestInferDeepPotLongRangeHPP, ValueTypes);
+
+TYPED_TEST(TestInferDeepPotLongRangeHPP, has_long_range) {
+  EXPECT_TRUE(this->dp.has_long_range());
+}
+
+TEST(TestInferDeepPotModelDeviLongRangeHPP, has_long_range) {
+#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT
+  GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#endif
+  // A long-range model anywhere in the ensemble confines all of them.
+  deepmd::hpp::DeepPotModelDevi mixed;
+  mixed.init(std::vector{kShortRangeModelPath, kModelPath});
+  EXPECT_TRUE(mixed.has_long_range());
+  deepmd::hpp::DeepPotModelDevi short_range;
+  short_range.init(
+      std::vector{kShortRangeModelPath, kShortRangeModelPath});
+  EXPECT_FALSE(short_range.has_long_range());
+}
+
+TYPED_TEST(TestInferDeepPotLongRangeHPP, nopbc) {
+  using VALUETYPE = TypeParam;
+  const double epsilon = std::is_same::value ? 1e-8 : 1e-4;
+  double ener;
+  std::vector force, virial;
+  this->dp.compute(ener, force, virial, this->coord, this->atype,
+                   std::vector());
+  ASSERT_EQ(force.size(), this->expected_f.size());
+  EXPECT_LT(fabs(ener - this->expected_tot_e), epsilon);
+  for (std::size_t ii = 0; ii < force.size(); ++ii) {
+    EXPECT_LT(fabs(force[ii] - this->expected_f[ii]), epsilon);
+  }
+}
diff --git a/source/api_c/tests/test_deeppot_model_devi_hpp.cc b/source/api_c/tests/test_deeppot_model_devi_hpp.cc
index 2a8242f99f..5fac31645d 100644
--- a/source/api_c/tests/test_deeppot_model_devi_hpp.cc
+++ b/source/api_c/tests/test_deeppot_model_devi_hpp.cc
@@ -120,6 +120,8 @@ TYPED_TEST(TestInferDeepPotModeDevi, attrs) {
   EXPECT_EQ(dp1.numb_types(), dp_md.numb_types());
   // EXPECT_EQ(dp1.dim_fparam(), dp_md.dim_fparam());
   // EXPECT_EQ(dp1.dim_aparam(), dp_md.dim_aparam());
+  EXPECT_FALSE(dp0.has_long_range());
+  EXPECT_FALSE(dp_md.has_long_range());
 }
 
 TYPED_TEST(TestInferDeepPotModeDevi, cpu_build_nlist) {
diff --git a/source/api_cc/CMakeLists.txt b/source/api_cc/CMakeLists.txt
index fe764ab0bf..a7a247ed58 100644
--- a/source/api_cc/CMakeLists.txt
+++ b/source/api_cc/CMakeLists.txt
@@ -121,6 +121,12 @@ if(ENABLE_PYTORCH AND "${OP_CXX_ABI_PT}" EQUAL "${OP_CXX_ABI}")
   deepmd_configure_backend_plugin(deepmd_backend_ptexpt)
   target_link_libraries(deepmd_backend_ptexpt PRIVATE "${TORCH_LIBRARIES}")
   target_compile_definitions(deepmd_backend_ptexpt PRIVATE BUILD_PYTORCH)
+  # Device-resident inference enqueues on the caller's stream, which needs the
+  # c10 CUDA stream API of a CUDA-enabled PyTorch.
+  if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA)
+    target_compile_definitions(deepmd_backend_ptexpt
+                               PRIVATE DEEPMD_TORCH_HAS_CUDA)
+  endif()
 endif()
 
 if(ENABLE_JAX)
diff --git a/source/api_cc/include/BackendPlugin.h b/source/api_cc/include/BackendPlugin.h
index 4149917414..ef92c4a464 100644
--- a/source/api_cc/include/BackendPlugin.h
+++ b/source/api_cc/include/BackendPlugin.h
@@ -15,15 +15,21 @@ class DipoleChargeModifierBase;
 
 /**
  * @brief C ABI symbol names exported by backend plugin libraries.
+ *
+ * The suffix of a symbol is the version of the interface it serves: that of
+ * the DeepPot and DeepSpin symbols changes with the virtual interface of
+ * DeepPotBackend and DeepSpinBackend, so that a plugin built against another
+ * interface is refused when it is loaded rather than called through a table
+ * of another layout.
  **/
 constexpr const char* DEEPMD_DEEPPOT_PLUGIN_CREATE_SYMBOL =
-    "deepmd_create_deeppot_backend_v1";
+    "deepmd_create_deeppot_backend_v2";
 constexpr const char* DEEPMD_DEEPPOT_PLUGIN_DELETE_SYMBOL =
-    "deepmd_delete_deeppot_backend_v1";
+    "deepmd_delete_deeppot_backend_v2";
 constexpr const char* DEEPMD_DEEPSPIN_PLUGIN_CREATE_SYMBOL =
-    "deepmd_create_deepspin_backend_v1";
+    "deepmd_create_deepspin_backend_v2";
 constexpr const char* DEEPMD_DEEPSPIN_PLUGIN_DELETE_SYMBOL =
-    "deepmd_delete_deepspin_backend_v1";
+    "deepmd_delete_deepspin_backend_v2";
 constexpr const char* DEEPMD_DEEPTENSOR_PLUGIN_CREATE_SYMBOL =
     "deepmd_create_deeptensor_backend_v1";
 constexpr const char* DEEPMD_DEEPTENSOR_PLUGIN_DELETE_SYMBOL =
diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h
index 575c7924c9..b1507d7103 100644
--- a/source/api_cc/include/DeepPot.h
+++ b/source/api_cc/include/DeepPot.h
@@ -7,6 +7,7 @@
 #include 
 
 #include "DeepBaseModel.h"
+#include "canonical_graph.h"
 #include "common.h"
 #include "neighbor_list.h"
 
@@ -424,19 +425,27 @@ class DeepPotBackend : public DeepBaseModelBackend {
   virtual void compute_canonical_graph_gpu(
       double* d_atom_energy,
       double* d_force,
+      double* d_virial,
       double* d_atom_virial,
       const std::int64_t* d_atype,
-      const std::uint32_t* d_source,
+      const void* d_source,
+      DP_CanonicalIndexWidth index_width,
       const float* d_edge_vec,
       const std::int64_t* d_destination_row_ptr,
       const std::int64_t* d_source_row_ptr,
-      const std::uint32_t* d_source_order,
-      const int nloc,
-      const int nall_nodes,
-      const std::int64_t edge_storage);
+      const std::int64_t* d_source_order,
+      const double* d_coord,
+      void* world,
+      const double* box,
+      const std::int64_t nloc,
+      const std::int64_t nall_nodes,
+      const std::int64_t edge_storage,
+      void* stream);
   virtual bool supports_device_edge_inference() const;
   virtual bool uses_fp32_edge_vectors() const;
   virtual bool uses_canonical_graph_inference() const;
+  virtual bool canonical_graph_requires_source_csr() const;
+  virtual bool has_long_range() const;
 };
 
 /**
@@ -894,18 +903,114 @@ class DeepPot : public DeepBaseModel {
    */
   bool uses_canonical_graph_inference() const;
 
+  /**
+   * @brief Whether every compact canonical evaluation needs the source CSR.
+   *
+   * An artifact whose force assembly groups edges by source reads the source
+   * CSR on every call. Any other artifact reads it only for the per-atom
+   * virial and accepts null source CSR pointers when that is not requested.
+   */
+  bool canonical_graph_requires_source_csr() const;
+
+  /**
+   * @brief Whether the model carries a long-range module.
+   *
+   * Such a module sums its interaction over every pair of atoms of the
+   * system, and over the images of the cell of a periodic system when it
+   * serves periodic systems, so the model evaluates complete systems, and
+   * compute_canonical_graph_gpu() passes it the positions of the nodes, the
+   * cell and, on several ranks, the communicator.
+   */
+  bool has_long_range() const;
+
+  /**
+   * @brief Evaluate a compact canonical graph from device-resident buffers.
+   *
+   * Every pointer resides on the model device.
+   *
+   * @param[out] d_atom_energy Per-local-node energy, shape ``(nloc)``.
+   * @param[out] d_force Per-node force, shape ``(nall_nodes, 3)``.
+   * @param[out] d_virial Virial of the graph, nine components, or null when
+   * it is not requested.
+   * @param[out] d_atom_virial Per-node virial, shape ``(nall_nodes, 9)``, or
+   * null when the per-atom virial is not requested.
+   * @param[in] d_atype Per-node atom types, shape ``(nall_nodes)``.
+   * @param[in] d_source Source node per edge storage slot.
+   * @param[in] index_width Width of each source index, four or eight bytes.
+   * @param[in] d_edge_vec FP32 edge vectors, shape ``(edge_storage, 3)``.
+   * @param[in] d_destination_row_ptr Destination CSR offsets.
+   * @param[in] d_source_row_ptr Source CSR offsets. May be null together with
+   * ``d_source_order`` and ``d_atom_virial`` unless
+   * canonical_graph_requires_source_csr().
+   * @param[in] d_source_order Source-grouped int64 edge storage positions.
+   * @param[in] d_coord Node positions in the length unit of the model, shape
+   * ``(nall_nodes, 3)``, in the order of ``d_atype``, or null. Only a model
+   * with a long-range module (has_long_range()) reads them, and it requires
+   * them; every other model ignores them.
+   * @param[in] world Pointer to the MPI communicator of the caller (an
+   * ``MPI_Comm``), as InputNlist::world, when the system spans several
+   * ranks, or null on a single rank. A model with a long-range module
+   * evaluates a complete system: a single rank holds every node of a
+   * non-periodic system (``nall_nodes == nloc``), and the halo nodes of a
+   * periodic one are images of its owned nodes; on several ranks it requires
+   * the communicator, reads the positions of the nodes each rank owns, and
+   * is evaluated on every rank at every step, a rank without nodes included
+   * (``nall_nodes``, ``nloc`` and ``edge_storage`` zero), whose buffers
+   * sized by the nodes or by the edge storage may then be null. Every other
+   * model ignores it.
+   * @param[in] box The cell of a periodic system in host memory, nine values
+   * in the length unit of the model, row-major with the lattice vectors as
+   * rows, or null for a non-periodic system. A model with a long-range
+   * module that serves periodic systems evaluates the system over the images
+   * of the cell, one whose module serves none rejects a cell, and every
+   * other model ignores it.
+   * @param[in] nloc Number of owned local nodes.
+   * @param[in] nall_nodes Total local-plus-halo node count.
+   * @param[in] edge_storage Number of edge storage slots.
+   * @param[in] stream Accelerator stream (``cudaStream_t``) of the model device
+   * on which the inputs were produced and the outputs will be consumed; a
+   * non-null stream requires a CUDA build. The evaluation is enqueued on it
+   * without a synchronization of the stream. A graph without halo nodes is
+   * captured into a CUDA graph on the second consecutive call with the same
+   * buffers and extents and replayed while they stay the same, except where a
+   * model with a long-range module receives a communicator. A model whose
+   * module serves periodic systems copies the cell into memory of its own,
+   * which a replay reads, and the mesh size of the cell counts as an extent: a
+   * call with another cell than the previous call first waits for the earlier
+   * evaluations, consecutive evaluations on different streams run in call
+   * order, and after the mesh size changes, the model evaluates without replay
+   * for 64 calls, after which the capture follows as for new extents. The host
+   * waits for the device when a call captures or releases a graph, which
+   * happens on that second call and on a call whose buffers or extents change.
+   * set_charge_spin waits for the work in flight on the model device before it
+   * replaces the constants and for the replacement before it returns, and the
+   * caller completes the work on its stream before it destroys the model. A
+   * null stream runs on the current stream of the calling thread for the model
+   * device and waits for that stream before returning. Models in different
+   * threads of one process evaluate concurrently, and their captures, releases
+   * of graphs and charge-state updates, which wait for the whole device, take
+   * turns; CUDA work outside DeePMD-kit that waits for the whole device while a
+   * capture is under way invalidates it, and the model then evaluates without
+   * replay.
+   */
   void compute_canonical_graph_gpu(double* d_atom_energy,
                                    double* d_force,
+                                   double* d_virial,
                                    double* d_atom_virial,
                                    const std::int64_t* d_atype,
-                                   const std::uint32_t* d_source,
+                                   const void* d_source,
+                                   DP_CanonicalIndexWidth index_width,
                                    const float* d_edge_vec,
                                    const std::int64_t* d_destination_row_ptr,
                                    const std::int64_t* d_source_row_ptr,
-                                   const std::uint32_t* d_source_order,
-                                   const int nloc,
-                                   const int nall_nodes,
-                                   const std::int64_t edge_storage);
+                                   const std::int64_t* d_source_order,
+                                   const double* d_coord,
+                                   void* world,
+                                   const double* box,
+                                   const std::int64_t nloc,
+                                   const std::int64_t nall_nodes,
+                                   const std::int64_t edge_storage,
+                                   void* stream = nullptr);
 
   int dim_chg_spin() const;
 
@@ -959,6 +1064,16 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
     return numb_models > 0 ? dps[0]->dim_chg_spin() : 0;
   };
 
+  /**
+   * @brief Whether any of the models carries a long-range module; see
+   * DeepPot::has_long_range().
+   **/
+  bool has_long_range() const {
+    return std::any_of(dps.begin(), dps.end(), [](const auto& model) {
+      return model->has_long_range();
+    });
+  };
+
   /**
    * @brief Fix the charge/spin condition served for the rest of the run.
    * Applied to every model, so that the deviation is taken between models
@@ -1134,5 +1249,16 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
 
  protected:
   std::vector> dps;
+
+ private:
+  /**
+   * @brief Check that every model accepts the cell handed to the ensemble.
+   * An empty cell, which a non-periodic system hands a model with a
+   * long-range module, is accepted only when every model carries one: the
+   * neighbor-list entries of other models, such as those of the TensorFlow
+   * backend, read the cell of the region.
+   **/
+  template 
+  void check_ensemble_box(const std::vector& dbox) const;
 };
 }  // namespace deepmd
diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h
index 4707ee9ce6..bc17f23c01 100644
--- a/source/api_cc/include/DeepPotPTExpt.h
+++ b/source/api_cc/include/DeepPotPTExpt.h
@@ -14,14 +14,21 @@
 
 #include 
 
+#include 
+#include 
+#include 
+#include 
+
 #include "DeepPot.h"
 #include "graph_assembly.h"
 
 // Forward-declare to keep these out of the public header. Defined in
-// commonPTExpt.h.
+// commonPTExpt.h, and PinnedCell in longRangePTExpt.h.
 namespace deepmd::ptexpt {
 class TempFile;
 class ChargeStateFold;
+class GraphReplay;
+class PinnedCell;
 }  // namespace deepmd::ptexpt
 
 namespace torch::inductor {
@@ -375,16 +382,22 @@ class DeepPotPTExpt : public DeepPotBackend {
                          const InputNlist* comm_nlist) override;
   void compute_canonical_graph_gpu(double* d_atom_energy,
                                    double* d_force,
+                                   double* d_virial,
                                    double* d_atom_virial,
                                    const std::int64_t* d_atype,
-                                   const std::uint32_t* d_source,
+                                   const void* d_source,
+                                   DP_CanonicalIndexWidth index_width,
                                    const float* d_edge_vec,
                                    const std::int64_t* d_destination_row_ptr,
                                    const std::int64_t* d_source_row_ptr,
-                                   const std::uint32_t* d_source_order,
-                                   const int nloc,
-                                   const int nall_nodes,
-                                   const std::int64_t edge_storage) override;
+                                   const std::int64_t* d_source_order,
+                                   const double* d_coord,
+                                   void* world,
+                                   const double* box,
+                                   const std::int64_t nloc,
+                                   const std::int64_t nall_nodes,
+                                   const std::int64_t edge_storage,
+                                   void* stream) override;
   /**
    * @brief FP32 edge-vector overload for compressed graph artifacts.
    *
@@ -413,6 +426,8 @@ class DeepPotPTExpt : public DeepPotBackend {
    */
   bool uses_fp32_edge_vectors() const override;
   bool uses_canonical_graph_inference() const override;
+  bool canonical_graph_requires_source_csr() const override;
+  bool has_long_range() const override;
 
  private:
   template 
@@ -432,16 +447,22 @@ class DeepPotPTExpt : public DeepPotBackend {
   void compute_canonical_graph_gpu_impl(
       double* d_atom_energy,
       double* d_force,
+      double* d_virial,
       double* d_atom_virial,
       const std::int64_t* d_atype,
-      const std::uint32_t* d_source,
+      const void* d_source,
+      DP_CanonicalIndexWidth index_width,
       const float* d_edge_vec,
       const std::int64_t* d_destination_row_ptr,
       const std::int64_t* d_source_row_ptr,
-      const std::uint32_t* d_source_order,
-      const int nloc,
-      const int nall_nodes,
-      const std::int64_t edge_storage);
+      const std::int64_t* d_source_order,
+      const double* d_coord,
+      void* world,
+      const double* box,
+      const std::int64_t nloc,
+      const std::int64_t nall_nodes,
+      const std::int64_t edge_storage,
+      void* stream);
   bool inited;
   // Every width below is a property of the loaded archive.  They read as zero
   // until ``init`` has run, so that a query on an uninitialised backend
@@ -472,7 +493,23 @@ class DeepPotPTExpt : public DeepPotBackend {
   bool lower_input_is_edge_ = false;
   bool lower_input_is_graph_ = false;
   bool lower_input_is_canonical_ = false;
+  // Whether the canonical force assembly groups edges by source, which makes
+  // the source CSR an input of every evaluation rather than of the per-atom
+  // virial alone.
+  bool canonical_requires_source_csr_ = false;
+  // Two-entry source CSR placeholders that decline the per-atom virial.
+  at::Tensor source_row_placeholder_;
+  at::Tensor source_order_placeholder_;
   bool graph_edge_fp32_ = false;
+  // The model carries a long-range module, whose lower reads the topologies
+  // of the frame after its other inputs: the pair lists, and the mesh of the
+  // cell of a module that serves periodic systems, for the graph lower; the
+  // node positions, the communicator of the ranks and the cell for the compact
+  // canonical lower.
+  bool long_range_ = false;
+  // The module serves periodic systems, on a mesh of at most this spacing.
+  bool long_range_periodic_ = false;
+  double mesh_spacing_ = 0.0;
   NeighborListData nlist_data;
   at::Tensor mapping_tensor;           // cached mapping tensor (LAMMPS path)
   std::vector mapping_;  // cached mapping vector (LAMMPS path)
@@ -480,6 +517,24 @@ class DeepPotPTExpt : public DeepPotBackend {
   at::Tensor edge_index_tensor;        // cached local edge graph (LAMMPS path)
   at::Tensor edge_index_ext_tensor;  // cached extended edge graph (LAMMPS path)
   std::unique_ptr loader;
+  // The cell input of the device-resident evaluations of a model with a
+  // long-range module, in pinned memory that a replay reads again.
+  std::unique_ptr long_range_cell_;
+  // The mesh size of the cell of the last device-resident evaluation of a
+  // model with a long-range module, zero for a non-periodic frame, and the
+  // evaluations since it last changed, counted up to the number of calls for
+  // which a changed mesh is evaluated eagerly before it is captured again.
+  std::optional > long_range_mesh_;
+  std::int64_t long_range_mesh_calls_ = 0;
+  // CUDA-graph replay of repeated device-resident canonical evaluations,
+  // declared after the loader and the cell so that the graph is released
+  // before the constants and the cell it reads.
+  std::unique_ptr canonical_replay_;
+  // Serializes the runs of the model, which runs on a single-threaded runner,
+  // the device-resident entry points, whose replay of canonical evaluations
+  // keeps its graph across calls, and the replacement of the constants. The
+  // mutex is recursive because a replayed evaluation enters the runner again.
+  std::recursive_mutex evaluation_mutex_;
   // Optional AOTInductor artifact for multi-rank message-passing inference,
   // loaded only when the .pt2 metadata declares ``has_comm_artifact``. It is
   // required when ghost features must be exchanged across ranks inside the
@@ -615,6 +670,11 @@ class DeepPotPTExpt : public DeepPotBackend {
    * @param[in] destination_row_ptr Destination CSR offsets ``(N + 1,)``.
    * @param[in] source_order Source-grouped edge permutation ``(E,)``.
    * @param[in] source_row_ptr Source CSR offsets ``(N + 1,)``.
+   * @param[in] coord Positions of the nodes of the single frame, shape
+   * ``(N, 3)`` float64, from which the topologies of a model with a
+   * long-range module are built; read only for such a model.
+   * @param[in] box The row-major cell of a periodic frame, nine values, or
+   * empty for a non-periodic one; read only for such a model.
    */
   std::vector run_model_graph(
       const torch::Tensor& atype,
@@ -629,8 +689,23 @@ class DeepPotPTExpt : public DeepPotBackend {
       const torch::Tensor& source_row_ptr,
       const torch::Tensor& fparam,
       const torch::Tensor& aparam,
-      const torch::Tensor& charge_spin);
+      const torch::Tensor& charge_spin,
+      const torch::Tensor& coord,
+      const std::vector& box);
 
+  /**
+   * @brief Run a compact canonical ``.pt2`` artifact.
+   *
+   * @param[in] coord Node positions, shape ``(N, 3)`` float64, appended for a
+   * model with a long-range module and ignored otherwise.
+   * @param[in] world The MPI communicator (a pointer to an ``MPI_Comm``) of the
+   * ranks a decomposed system spans, or NULL for a system that one rank holds
+   * whole; a model with a long-range module receives it after the positions.
+   * @param[in] cell The cell, a float64 CPU tensor of shape ``(1, 9)``: the
+   * row-major cell of a periodic frame, or zero for a non-periodic one; a
+   * model with a long-range module receives it after the communicator, and
+   * every other model ignores it.
+   */
   std::vector run_model_canonical_graph(
       const torch::Tensor& atype,
       const torch::Tensor& n_node,
@@ -639,7 +714,30 @@ class DeepPotPTExpt : public DeepPotBackend {
       const torch::Tensor& edge_vec,
       const torch::Tensor& destination_row_ptr,
       const torch::Tensor& source_row_ptr,
-      const torch::Tensor& source_order);
+      const torch::Tensor& source_order,
+      const torch::Tensor& coord,
+      const void* world,
+      const torch::Tensor& cell);
+
+  /**
+   * @brief Evaluate the long-range module on a rank that holds no node.
+   *
+   * The compact lower runs on one halo node without edges, which carries no
+   * charge and contributes nothing, so that the rank joins the exchanges of
+   * the module with the other ranks of ``world``, in the cell ``box`` of the
+   * system, empty for a non-periodic one.
+   */
+  void join_long_range_exchange(const void* world,
+                                const std::vector& box);
+
+  /**
+   * @brief Check the cell given to a model with a long-range module.
+   *
+   * @param[in] box The row-major cell of a periodic frame, nine values, or
+   * empty for a non-periodic one, which every long-range module serves; a
+   * cell requires a module that serves periodic systems.
+   */
+  void check_long_range_box(const std::vector& box) const;
 
   /**
    * @brief Run the with-comm .pt2 artifact with comm tensors appended.
diff --git a/source/api_cc/include/DeepSpin.h b/source/api_cc/include/DeepSpin.h
index bf57849b14..338d76ff85 100644
--- a/source/api_cc/include/DeepSpin.h
+++ b/source/api_cc/include/DeepSpin.h
@@ -5,6 +5,7 @@
 #include 
 
 #include "DeepBaseModel.h"
+#include "canonical_graph.h"
 #include "common.h"
 #include "neighbor_list.h"
 
@@ -290,17 +291,20 @@ class DeepSpinBackend : public DeepBaseModelBackend {
       double* d_atom_energy,
       double* d_force,
       double* d_force_mag,
+      double* d_virial,
       double* d_atom_virial,
       const std::int64_t* d_atype,
-      const std::uint32_t* d_source,
+      const void* d_source,
+      DP_CanonicalIndexWidth index_width,
       const float* d_edge_vec,
       const std::int64_t* d_destination_row_ptr,
       const std::int64_t* d_source_row_ptr,
-      const std::uint32_t* d_source_order,
+      const std::int64_t* d_source_order,
       const float* d_spin,
-      const int nloc,
-      const int nall_nodes,
-      const std::int64_t edge_storage);
+      const std::int64_t nloc,
+      const std::int64_t nall_nodes,
+      const std::int64_t edge_storage,
+      void* stream);
   virtual bool uses_canonical_graph_inference() const;
 
   /**
@@ -634,33 +638,44 @@ class DeepSpin : public DeepBaseModel {
    * @param[out] d_force Per-node force, [nall_nodes * 3] row-major.
    * @param[out] d_force_mag Per-node magnetic force, [nall_nodes * 3]
    *row-major.
-   * @param[out] d_atom_virial Per-node virial, [nall_nodes * 9] row-major.
+   * @param[out] d_virial Virial of the graph, [9], or null when it is not
+   *requested.
+   * @param[out] d_atom_virial Per-node virial, [nall_nodes * 9] row-major, or
+   *null when the per-atom virial is not requested.
    * @param[in] d_atype Per-node atom types, [nall_nodes].
    * @param[in] d_source Source-node index per edge, [edge_storage].
    * @param[in] d_edge_vec Destination-major edge vectors, [edge_storage * 3].
    * @param[in] d_destination_row_ptr Destination CSR offsets, [nall_nodes + 1].
-   * @param[in] d_source_row_ptr Source CSR offsets, [nall_nodes + 1].
+   * @param[in] d_source_row_ptr Source CSR offsets, [nall_nodes + 1]. The
+   *source CSR is read only for the per-atom virial and may be null, together
+   *with ``d_source_order``, when ``d_atom_virial`` is null.
    * @param[in] d_source_order Source-grouped edge positions, [edge_storage].
    * @param[in] d_spin Per-node magnetic moment, [nall_nodes * 3] row-major;
    *ghost rows carry their owner's moment.
    * @param[in] nloc Number of local atoms.
    * @param[in] nall_nodes Graph node count (local + ghost).
    * @param[in] edge_storage Allocated edge capacity.
+   * @param[in] stream Accelerator stream (``cudaStream_t``) of the model
+   *device on which the inputs were produced and the outputs will be consumed,
+   *with the semantics of DeepPot::compute_canonical_graph_gpu.
    */
   void compute_canonical_graph_gpu(double* d_atom_energy,
                                    double* d_force,
                                    double* d_force_mag,
+                                   double* d_virial,
                                    double* d_atom_virial,
                                    const std::int64_t* d_atype,
-                                   const std::uint32_t* d_source,
+                                   const void* d_source,
+                                   DP_CanonicalIndexWidth index_width,
                                    const float* d_edge_vec,
                                    const std::int64_t* d_destination_row_ptr,
                                    const std::int64_t* d_source_row_ptr,
-                                   const std::uint32_t* d_source_order,
+                                   const std::int64_t* d_source_order,
                                    const float* d_spin,
-                                   const int nloc,
-                                   const int nall_nodes,
-                                   const std::int64_t edge_storage);
+                                   const std::int64_t nloc,
+                                   const std::int64_t nall_nodes,
+                                   const std::int64_t edge_storage,
+                                   void* stream = nullptr);
 
  protected:
   std::shared_ptr dp;
diff --git a/source/api_cc/include/NativeSpinPTExpt.h b/source/api_cc/include/NativeSpinPTExpt.h
index bb8f820980..c7516ceab4 100644
--- a/source/api_cc/include/NativeSpinPTExpt.h
+++ b/source/api_cc/include/NativeSpinPTExpt.h
@@ -14,13 +14,16 @@
 
 #include 
 
+#include 
+
 #include "DeepSpin.h"
 
 // Forward-declare to keep the private header out of the public one. Defined in
 // commonPTExpt.h.
 namespace deepmd::ptexpt {
 class ChargeStateFold;
-}
+class GraphReplay;
+}  // namespace deepmd::ptexpt
 
 namespace torch::inductor {
 class AOTIModelPackageLoader;
@@ -52,8 +55,9 @@ struct CanonicalGraphTensorPack;
  *   graph-lower descriptor can be frozen this way.
  * - ``lower_input_kind == "dpa4c_canonical"``, the compact deployment ABI
  *   (``forward_lower_canonical_graph_exportable``): the eight dual-CSR graph
- *   tensors -- uint32 topology and float32 edge vectors -- and the moment at
- *   positional index 8, with no conditional tail.
+ *   tensors -- source words of four or eight bytes, int64 row pointers and
+ *   positions, and float32 edge vectors -- and the moment at positional
+ *   index 8, with no conditional tail.
  *
  * Three entry points share the selected forward:
  *
@@ -289,7 +293,10 @@ class NativeSpinPTExpt : public DeepSpinBackend {
    * @param[out] d_force Per-node force, GPU [nall_nodes * 3] row-major.
    * @param[out] d_force_mag Per-node magnetic force, GPU [nall_nodes * 3]
    *   row-major.
-   * @param[out] d_atom_virial Per-node virial, GPU [nall_nodes * 9] row-major.
+   * @param[out] d_virial Virial of the graph, GPU [9], or null when it is not
+   *   requested.
+   * @param[out] d_atom_virial Per-node virial, GPU [nall_nodes * 9] row-major,
+   *   or null when the per-atom virial is not requested.
    * @param[in] d_atype Per-node atom types, GPU [nall_nodes].
    * @param[in] d_source Source-node index per edge, GPU [edge_storage].
    * @param[in] d_edge_vec Destination-major edge vectors, GPU
@@ -297,6 +304,8 @@ class NativeSpinPTExpt : public DeepSpinBackend {
    * @param[in] d_destination_row_ptr Destination CSR offsets, GPU
    *   [nall_nodes + 1].
    * @param[in] d_source_row_ptr Source CSR offsets, GPU [nall_nodes + 1].
+   *   Read only for the per-atom virial; may be null, together with
+   *   ``d_source_order``, when ``d_atom_virial`` is null.
    * @param[in] d_source_order Source-grouped edge positions, GPU
    *   [edge_storage].
    * @param[in] d_spin Per-node magnetic moment, GPU [nall_nodes * 3]
@@ -304,21 +313,25 @@ class NativeSpinPTExpt : public DeepSpinBackend {
    * @param[in] nloc Number of local atoms.
    * @param[in] nall_nodes Graph node count (local + ghost).
    * @param[in] edge_storage Allocated edge capacity.
+   * @param[in] stream Caller stream; see DeepSpin::compute_canonical_graph_gpu.
    */
   void compute_canonical_graph_gpu(double* d_atom_energy,
                                    double* d_force,
                                    double* d_force_mag,
+                                   double* d_virial,
                                    double* d_atom_virial,
                                    const std::int64_t* d_atype,
-                                   const std::uint32_t* d_source,
+                                   const void* d_source,
+                                   DP_CanonicalIndexWidth index_width,
                                    const float* d_edge_vec,
                                    const std::int64_t* d_destination_row_ptr,
                                    const std::int64_t* d_source_row_ptr,
-                                   const std::uint32_t* d_source_order,
+                                   const std::int64_t* d_source_order,
                                    const float* d_spin,
-                                   const int nloc,
-                                   const int nall_nodes,
-                                   const std::int64_t edge_storage) override;
+                                   const std::int64_t nloc,
+                                   const std::int64_t nall_nodes,
+                                   const std::int64_t edge_storage,
+                                   void* stream) override;
 
   bool uses_canonical_graph_inference() const override;
 
@@ -502,6 +515,9 @@ class NativeSpinPTExpt : public DeepSpinBackend {
   // Which of the two schemas the loaded artifact declares: the compact
   // dual-CSR ABI when true, the general NeighborGraph ABI when false.
   bool canonical_abi_ = false;
+  // Two-entry source CSR placeholders that decline the per-atom virial.
+  at::Tensor source_row_placeholder_;
+  at::Tensor source_order_placeholder_;
   // Edge-vector precision the NeighborGraph artifact was traced with, read
   // from the ``graph_edge_dtype`` metadata field. The compact ABI is float32
   // by definition and validates that at load.
@@ -525,6 +541,15 @@ class NativeSpinPTExpt : public DeepSpinBackend {
   at::Tensor edge_index_tensor;      // node-space edges (folded or extended)
   at::Tensor edge_index_ext_tensor;  // extended-atom edges, for the geometry
   std::unique_ptr loader;
+  // CUDA-graph replay of repeated device-resident canonical evaluations,
+  // declared after the loader so that the graph is released before the
+  // constants it reads.
+  std::unique_ptr canonical_replay_;
+  // Serializes the runs of the model, which runs on a single-threaded runner,
+  // the device-resident entry point, whose replay of canonical evaluations
+  // keeps its graph across calls, and the replacement of the constants. The
+  // mutex is recursive because a replayed evaluation enters the runner again.
+  std::recursive_mutex evaluation_mutex_;
   // The charge/spin condition a compressed descriptor folded into the
   // constants of its lower, re-runnable so that a condition chosen at runtime
   // can be served.  Null for an uncompressed model, which reads its condition
diff --git a/source/api_cc/include/canonical_graph.h b/source/api_cc/include/canonical_graph.h
new file mode 100644
index 0000000000..359569ec31
--- /dev/null
+++ b/source/api_cc/include/canonical_graph.h
@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+#pragma once
+
+#include 
+
+/** Width of node indices in bytes; edge counts and CSR positions are int64. */
+typedef enum DP_CanonicalIndexWidth {
+  DP_CANONICAL_INDEX32 = 4,
+  DP_CANONICAL_INDEX64 = 8
+} DP_CanonicalIndexWidth;
+
+#ifdef __cplusplus
+namespace deepmd {
+/** Select node-index width, preserving signed words at native host boundaries.
+ */
+inline DP_CanonicalIndexWidth canonical_index_width(int64_t node_count) {
+  return node_count <= (int64_t{1} << 31) ? DP_CANONICAL_INDEX32
+                                          : DP_CANONICAL_INDEX64;
+}
+}  // namespace deepmd
+#endif
diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h
index 873158615c..999b3670a7 100644
--- a/source/api_cc/include/commonPT.h
+++ b/source/api_cc/include/commonPT.h
@@ -16,6 +16,7 @@
 #include 
 #include 
 
+#include "canonical_graph.h"
 #include "common.h"
 #include "neighbor_list.h"
 
@@ -418,21 +419,12 @@ struct CanonicalGraphTensorPack {
   torch::Tensor source_order;
 };
 
-/**
- * @brief Return the scalar type of compact canonical graph indices.
- *
- * Unsigned 32-bit tensors entered the public C++ API in PyTorch 2.3. Older
- * CPU-only libtorch releases can still build DeePMD-kit, but cannot execute
- * the GPU-only compact canonical graph path.
+/** Expose typed canonical indices as a zero-copy byte carrier for AOTInductor.
  */
-inline at::ScalarType canonicalGraphIndexType() {
-#if TORCH_VERSION_MAJOR > 2 || \
-    (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 3)
-  return torch::kUInt32;
-#else
-  throw deepmd_exception(
-      "compact canonical graph inference requires PyTorch 2.3 or later");
-#endif
+inline torch::Tensor canonicalIndexBytes(const torch::Tensor& index) {
+  auto shape = index.sizes().vec();
+  shape.push_back(index.element_size());
+  return index.view(torch::kUInt8).view(shape);
 }
 
 inline CanonicalGraphTensorPack compactCanonicalGraph(
@@ -441,20 +433,16 @@ inline CanonicalGraphTensorPack compactCanonicalGraph(
       graph.destination_row_ptr.select(0, graph.destination_row_ptr.size(0) - 1)
           .item();
   const std::int64_t storage_count = std::max(edge_count, 2);
-  if (static_cast(storage_count) >
-      std::numeric_limits::max()) {
-    throw deepmd_exception(
-        "compact canonical graph exceeds the uint32 edge-index range");
-  }
-  const auto index_type = canonicalGraphIndexType();
+  const auto index_type =
+      canonical_index_width(graph.atype.size(0)) == DP_CANONICAL_INDEX32
+          ? torch::kInt32
+          : torch::kInt64;
   auto source = torch::zeros({storage_count},
                              graph.edge_index.options().dtype(index_type));
   auto edge_vec = torch::zeros({storage_count, 3},
                                graph.edge_vec.options().dtype(torch::kFloat32));
-  auto source_order =
-      torch::arange(storage_count,
-                    graph.edge_index.options().dtype(torch::kInt64))
-          .to(index_type);
+  auto source_order = torch::arange(
+      storage_count, graph.edge_index.options().dtype(torch::kInt64));
   if (edge_count > 0) {
     source.slice(0, 0, edge_count)
         .copy_(graph.edge_index.select(0, 0)
@@ -463,7 +451,7 @@ inline CanonicalGraphTensorPack compactCanonicalGraph(
     edge_vec.slice(0, 0, edge_count)
         .copy_(graph.edge_vec.slice(0, 0, edge_count).to(torch::kFloat32));
     source_order.slice(0, 0, edge_count)
-        .copy_(graph.source_order.slice(0, 0, edge_count).to(index_type));
+        .copy_(graph.source_order.slice(0, 0, edge_count));
   }
   return {graph.atype,
           graph.n_node,
diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc
index 3373d6fe71..68de144850 100644
--- a/source/api_cc/src/DeepPot.cc
+++ b/source/api_cc/src/DeepPot.cc
@@ -659,18 +659,25 @@ void DeepPotBackend::compute_edges_gpu(double* d_atom_energy,
 void DeepPotBackend::compute_canonical_graph_gpu(
     double* d_atom_energy,
     double* d_force,
+    double* d_virial,
     double* d_atom_virial,
     const std::int64_t* d_atype,
-    const std::uint32_t* d_source,
+    const void* d_source,
+    DP_CanonicalIndexWidth index_width,
     const float* d_edge_vec,
     const std::int64_t* d_destination_row_ptr,
     const std::int64_t* d_source_row_ptr,
-    const std::uint32_t* d_source_order,
-    const int nloc,
-    const int nall_nodes,
-    const std::int64_t edge_storage) {
+    const std::int64_t* d_source_order,
+    const double* d_coord,
+    void* world,
+    const double* box,
+    const std::int64_t nloc,
+    const std::int64_t nall_nodes,
+    const std::int64_t edge_storage,
+    void* stream) {
   (void)d_atom_energy;
   (void)d_force;
+  (void)d_virial;
   (void)d_atom_virial;
   (void)d_atype;
   (void)d_source;
@@ -678,9 +685,13 @@ void DeepPotBackend::compute_canonical_graph_gpu(
   (void)d_destination_row_ptr;
   (void)d_source_row_ptr;
   (void)d_source_order;
+  (void)d_coord;
+  (void)world;
+  (void)box;
   (void)nloc;
   (void)nall_nodes;
   (void)edge_storage;
+  (void)stream;
   throw deepmd::deepmd_exception(
       "compact canonical graph inference is only supported by a compatible "
       "PyTorch Exportable backend.");
@@ -688,10 +699,16 @@ void DeepPotBackend::compute_canonical_graph_gpu(
 
 bool DeepPotBackend::uses_fp32_edge_vectors() const { return false; }
 
+bool DeepPotBackend::canonical_graph_requires_source_csr() const {
+  return false;
+}
+
 bool DeepPotBackend::supports_device_edge_inference() const { return false; }
 
 bool DeepPotBackend::uses_canonical_graph_inference() const { return false; }
 
+bool DeepPotBackend::has_long_range() const { return false; }
+
 void DeepPot::compute_edges_gpu(double* d_atom_energy,
                                 double* d_force,
                                 double* d_atom_virial,
@@ -748,20 +765,27 @@ void DeepPot::compute_edges_gpu(double* d_atom_energy,
 void DeepPot::compute_canonical_graph_gpu(
     double* d_atom_energy,
     double* d_force,
+    double* d_virial,
     double* d_atom_virial,
     const std::int64_t* d_atype,
-    const std::uint32_t* d_source,
+    const void* d_source,
+    DP_CanonicalIndexWidth index_width,
     const float* d_edge_vec,
     const std::int64_t* d_destination_row_ptr,
     const std::int64_t* d_source_row_ptr,
-    const std::uint32_t* d_source_order,
-    const int nloc,
-    const int nall_nodes,
-    const std::int64_t edge_storage) {
+    const std::int64_t* d_source_order,
+    const double* d_coord,
+    void* world,
+    const double* box,
+    const std::int64_t nloc,
+    const std::int64_t nall_nodes,
+    const std::int64_t edge_storage,
+    void* stream) {
   dp->compute_canonical_graph_gpu(
-      d_atom_energy, d_force, d_atom_virial, d_atype, d_source, d_edge_vec,
-      d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes,
-      edge_storage);
+      d_atom_energy, d_force, d_virial, d_atom_virial, d_atype, d_source,
+      index_width, d_edge_vec, d_destination_row_ptr, d_source_row_ptr,
+      d_source_order, d_coord, world, box, nloc, nall_nodes, edge_storage,
+      stream);
 }
 
 bool DeepPot::uses_fp32_edge_vectors() const {
@@ -776,6 +800,12 @@ bool DeepPot::uses_canonical_graph_inference() const {
   return dp->uses_canonical_graph_inference();
 }
 
+bool DeepPot::canonical_graph_requires_source_csr() const {
+  return dp->canonical_graph_requires_source_csr();
+}
+
+bool DeepPot::has_long_range() const { return dp->has_long_range(); }
+
 int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); }
 
 void DeepPot::set_charge_spin(const std::vector& charge_spin) {
@@ -921,6 +951,25 @@ template void DeepPotModelDevi::compute(
     const std::vector& aparam,
     const std::vector& charge_spin);
 
+// LAMMPS hands a model with a long-range module no cell for a non-periodic
+// system, whereas the neighbor-list entries of other models, such as those of
+// the TensorFlow backend, read the cell of the region.
+template 
+void DeepPotModelDevi::check_ensemble_box(
+    const std::vector& dbox) const {
+  if (!dbox.empty() || !has_long_range()) {
+    return;
+  }
+  for (const auto& model : dps) {
+    if (!model->has_long_range()) {
+      throw deepmd::deepmd_exception(
+          "an ensemble that holds a model with a long-range module evaluates "
+          "a non-periodic system only when every model carries one; use "
+          "models with long-range modules throughout, or a periodic system.");
+    }
+  }
+}
+
 template 
 void DeepPotModelDevi::compute(std::vector& all_energy,
                                std::vector>& all_force,
@@ -937,6 +986,7 @@ void DeepPotModelDevi::compute(std::vector& all_energy,
   if (numb_models == 0) {
     return;
   }
+  check_ensemble_box(dbox);
   all_energy.resize(numb_models);
   all_force.resize(numb_models);
   all_virial.resize(numb_models);
@@ -999,6 +1049,7 @@ void DeepPotModelDevi::compute(
   all_virial.resize(numb_models);
   all_atom_energy.resize(numb_models);
   all_atom_virial.resize(numb_models);
+  check_ensemble_box(dbox);
   for (unsigned ii = 0; ii < numb_models; ++ii) {
     dps[ii]->compute(all_energy[ii], all_force[ii], all_virial[ii],
                      all_atom_energy[ii], all_atom_virial[ii], dcoord_, datype_,
diff --git a/source/api_cc/src/DeepPotJAXPlugin.cc b/source/api_cc/src/DeepPotJAXPlugin.cc
index ac2d5485a1..2522148833 100644
--- a/source/api_cc/src/DeepPotJAXPlugin.cc
+++ b/source/api_cc/src/DeepPotJAXPlugin.cc
@@ -4,7 +4,7 @@
 #include "BackendPluginFactory.h"
 #include "DeepPotJAX.h"
 
-extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
+extern "C" void* deepmd_create_deeppot_backend_v2(const char* model,
                                                   int gpu_rank,
                                                   const char* file_content,
                                                   std::size_t file_content_size,
@@ -13,7 +13,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
       model, gpu_rank, file_content, file_content_size, error_message);
 }
 
-extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) {
+extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) {
   deepmd::plugin::delete_deeppot_backend(backend);
 }
 
diff --git a/source/api_cc/src/DeepPotPDPlugin.cc b/source/api_cc/src/DeepPotPDPlugin.cc
index fbc920cfb5..3bbd6842de 100644
--- a/source/api_cc/src/DeepPotPDPlugin.cc
+++ b/source/api_cc/src/DeepPotPDPlugin.cc
@@ -4,7 +4,7 @@
 #include "BackendPluginFactory.h"
 #include "DeepPotPD.h"
 
-extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
+extern "C" void* deepmd_create_deeppot_backend_v2(const char* model,
                                                   int gpu_rank,
                                                   const char* file_content,
                                                   std::size_t file_content_size,
@@ -13,7 +13,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
       model, gpu_rank, file_content, file_content_size, error_message);
 }
 
-extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) {
+extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) {
   deepmd::plugin::delete_deeppot_backend(backend);
 }
 
diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc
index 005a18e722..c92252c1e2 100644
--- a/source/api_cc/src/DeepPotPTExpt.cc
+++ b/source/api_cc/src/DeepPotPTExpt.cc
@@ -7,6 +7,7 @@
 #include 
 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -22,6 +23,7 @@
 #include "commonPTExpt.h"
 #include "device.h"
 #include "errors.h"
+#include "longRangePTExpt.h"
 #include "neighbor_list.h"
 
 using deepmd::ptexpt::check_call_charge_spin;
@@ -38,10 +40,19 @@ namespace {
 // commonPT.h (deepmd namespace) so DeepSpinPTExpt.cc can share them for its
 // native-spin graph route.
 
-void synchronize_current_accelerator_stream() {
-#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
+// Wait for the evaluation that the calling thread enqueued on its current
+// stream of ``device``, on which the device-resident entry points without a
+// caller's stream enqueue all of their work; the work and the CUDA graph
+// captures of other threads on the device proceed meanwhile.
+void synchronize_current_accelerator_stream(const int device) {
+#if defined(DEEPMD_TORCH_HAS_CUDA)
+  c10::cuda::getCurrentCUDAStream(static_cast(device))
+      .synchronize();
+#elif defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
+  (void)device;
   DPErrcheck(gpuDeviceSynchronize());
 #else
+  (void)device;
   throw deepmd::deepmd_exception(
       "GPU-resident inference requires a GPU-enabled DeePMD-kit build.");
 #endif
@@ -110,6 +121,37 @@ FrameParameterLayout resolve_frame_parameter_layout(const char* name,
       " for frame-specific values.");
 }
 
+// The communicator input of the compact lower of a long-range model: an int64
+// CPU tensor holding the address of the ``MPI_Comm`` that ``world`` points to,
+// or zero for a system that one rank holds whole. The handle itself is read
+// only by the operator, with the width of ``MPI_Comm`` in its MPI library.
+at::Tensor communicator_tensor(const void* world) {
+  return torch::full(
+      {1}, static_cast(reinterpret_cast(world)),
+      torch::TensorOptions().dtype(torch::kInt64));
+}
+
+// A capture waits for the device, empties the cache of the allocator and
+// allocates the memory of its graph anew, which takes far longer than an
+// evaluation of a small system. After the mesh of its cell changes size, a
+// model with a long-range module is therefore evaluated eagerly until the
+// size has held for this many calls, so that a cell that fluctuates across a
+// boundary of the mesh size under a barostat is not captured at every
+// crossing.
+constexpr std::int64_t kStableMeshCalls = 64;
+
+// The cell input of the compact lower of a long-range model: a float64 CPU
+// tensor of shape (1, 9), the rows of the row-major cell its lattice vectors,
+// zero for a non-periodic frame.
+at::Tensor cell_tensor(const std::vector& box) {
+  auto cell =
+      torch::zeros({1, 9}, torch::TensorOptions().dtype(torch::kFloat64));
+  if (!box.empty()) {
+    std::copy(box.begin(), box.begin() + 9, cell.data_ptr());
+  }
+  return cell;
+}
+
 }  // namespace
 
 void DeepPotPTExpt::translate_error(std::function f) {
@@ -244,6 +286,7 @@ void DeepPotPTExpt::init(const std::string& model,
       nnei += v.as_int();
     }
   }
+  long_range_ = metadata.obj_val.count("topologies") > 0;
   if (metadata.obj_val.count("lower_input_kind")) {
     const std::string lower_input_kind =
         metadata["lower_input_kind"].as_string();
@@ -251,10 +294,43 @@ void DeepPotPTExpt::init(const std::string& model,
     lower_input_is_graph_ = lower_input_kind == "graph";
     lower_input_is_canonical_ = lower_input_kind == "dpa1_canonical" ||
                                 lower_input_kind == "dpa4c_canonical";
+    // Whether the force assembly of the compact lower reads the source CSR on
+    // every call, as the archive records it; an archive without the record
+    // reads it for compressed DPA1 and for a long-range module.
+    canonical_requires_source_csr_ =
+        metadata.obj_val.count("canonical_requires_source_csr")
+            ? metadata["canonical_requires_source_csr"].as_bool()
+            : lower_input_kind == "dpa1_canonical" ||
+                  (lower_input_is_canonical_ && long_range_);
   } else {
     lower_input_is_edge_ = false;
     lower_input_is_graph_ = false;
     lower_input_is_canonical_ = false;
+    canonical_requires_source_csr_ = false;
+  }
+  // A model with a long-range module reads the all-pairs topology of the
+  // frame after its other inputs, and, when the module serves periodic
+  // systems, the mesh of the cell, sized by the spacing of the metadata.
+  long_range_periodic_ = false;
+  if (long_range_) {
+    for (const auto& name : metadata["topologies"].as_array()) {
+      if (name.as_string() == "pme_mesh") {
+        long_range_periodic_ = true;
+      } else if (name.as_string() != "all_pairs") {
+        throw deepmd::deepmd_exception("unknown topology '" + name.as_string() +
+                                       "' in the .pt2 metadata; expected "
+                                       "'all_pairs' or 'pme_mesh'.");
+      }
+    }
+    if (long_range_periodic_) {
+      mesh_spacing_ =
+          metadata["topology_params"]["pme_mesh"]["spacing"].as_double();
+    }
+    if (!lower_input_is_graph_ && !lower_input_is_canonical_) {
+      throw deepmd::deepmd_exception(
+          "a model with a long-range module requires the graph or the "
+          "compact canonical lower.");
+    }
   }
   if (lower_input_is_edge_) {
     std::cerr << "WARNING: This .pt2 uses the deprecated edge_vec lower "
@@ -279,9 +355,12 @@ void DeepPotPTExpt::init(const std::string& model,
           "compact canonical graph artifacts require float32 edge vectors.");
     }
     if (!metadata.obj_val.count("canonical_index_dtype") ||
-        metadata["canonical_index_dtype"].as_string() != "uint32") {
+        metadata["canonical_index_dtype"].as_string() != "auto" ||
+        !metadata.obj_val.count("canonical_index_storage") ||
+        metadata["canonical_index_storage"].as_string() != "byte_words") {
       throw deepmd::deepmd_exception(
-          "compact canonical graph artifacts require uint32 topology; "
+          "compact canonical graph artifacts require automatic-width byte "
+          "indices; "
           "re-freeze the model with the current DeePMD-kit version.");
     }
   }
@@ -297,9 +376,13 @@ void DeepPotPTExpt::init(const std::string& model,
     output_keys.push_back(v.as_string());
   }
 
-  // Load the AOTInductor model package
+  // Load the AOTInductor model package. An instance runs its model one call
+  // at a time (evaluation_mutex_), so the model runs on the single-threaded
+  // runner, which records no completion event per run: a run is free to be
+  // enqueued before the previous one finishes and to be captured into a CUDA
+  // graph.
   loader = std::make_unique(
-      model, "model", false, 1,
+      model, "model", /*run_single_threaded=*/true, 1,
       gpu_enabled ? static_cast(gpu_id)
                   : static_cast(-1));
 
@@ -445,10 +528,27 @@ void DeepPotPTExpt::set_charge_spin(const std::vector& charge_spin) {
   // constants without opening a second failure point.
   std::vector next_charge_spin(charge_spin);
   if (charge_state_fold_) {
-    charge_state_fold_->apply(charge_spin,
-                              gpu_enabled ? torch::Device(torch::kCUDA, gpu_id)
-                                          : torch::Device(torch::kCPU),
-                              *loader);
+    const std::lock_guard evaluation(evaluation_mutex_);
+    translate_error([&] {
+      // Evaluations still in flight on a caller's stream, and a captured
+      // evaluation, read the constants the fold replaces, and the next
+      // evaluation, possibly on another stream, reads those it writes; the
+      // fold therefore waits for the device on both sides.
+      if (gpu_enabled) {
+        ptexpt::synchronize_device(gpu_id);
+      }
+      if (canonical_replay_) {
+        canonical_replay_->invalidate();
+      }
+      charge_state_fold_->apply(charge_spin,
+                                gpu_enabled
+                                    ? torch::Device(torch::kCUDA, gpu_id)
+                                    : torch::Device(torch::kCPU),
+                                *loader);
+      if (gpu_enabled) {
+        ptexpt::synchronize_device(gpu_id);
+      }
+    });
   }
   default_chg_spin_.swap(next_charge_spin);
 }
@@ -473,6 +573,7 @@ std::vector DeepPotPTExpt::run_model(
   if (dchgspin > 0) {
     inputs.push_back(charge_spin);
   }
+  const std::lock_guard evaluation(evaluation_mutex_);
   return loader->run(inputs);
 }
 
@@ -497,6 +598,7 @@ std::vector DeepPotPTExpt::run_model_edges(
   if (dchgspin > 0) {
     inputs.push_back(charge_spin);
   }
+  const std::lock_guard evaluation(evaluation_mutex_);
   return loader->run(inputs);
 }
 
@@ -513,7 +615,9 @@ std::vector DeepPotPTExpt::run_model_graph(
     const torch::Tensor& source_row_ptr,
     const torch::Tensor& fparam,
     const torch::Tensor& aparam,
-    const torch::Tensor& charge_spin) {
+    const torch::Tensor& charge_spin,
+    const torch::Tensor& coord,
+    const std::vector& box) {
   // Graph-input ABI: original edge payload plus destination/source CSR views.
   check_graph_aparam_flat(aparam, daparam, "run_model_graph");
   std::vector inputs = {
@@ -529,6 +633,22 @@ std::vector DeepPotPTExpt::run_model_graph(
   if (dchgspin > 0) {
     inputs.push_back(charge_spin);
   }
+  if (long_range_) {
+    // A periodic frame holds no pair; the mesh describes its interactions.
+    const auto vec_dtype = graph_edge_fp32_ ? torch::kFloat32 : torch::kFloat64;
+    const auto pairs =
+        box.empty()
+            ? deepmd::ptexpt::buildAllPairsTopology(
+                  coord, atype, pair_exclude_table_, ntypes, vec_dtype)
+            : deepmd::ptexpt::buildGuardPairTopology(coord.device(), vec_dtype);
+    inputs.insert(inputs.end(), pairs.begin(), pairs.end());
+    if (long_range_periodic_) {
+      const auto mesh = deepmd::ptexpt::buildMeshTopology(
+          coord, box, mesh_spacing_, vec_dtype);
+      inputs.insert(inputs.end(), mesh.begin(), mesh.end());
+    }
+  }
+  const std::lock_guard evaluation(evaluation_mutex_);
   return loader->run(inputs);
 }
 
@@ -540,9 +660,72 @@ std::vector DeepPotPTExpt::run_model_canonical_graph(
     const torch::Tensor& edge_vec,
     const torch::Tensor& destination_row_ptr,
     const torch::Tensor& source_row_ptr,
-    const torch::Tensor& source_order) {
-  return loader->run({atype, n_node, n_local, source, edge_vec,
-                      destination_row_ptr, source_row_ptr, source_order});
+    const torch::Tensor& source_order,
+    const torch::Tensor& coord,
+    const void* world,
+    const torch::Tensor& cell) {
+  std::vector inputs = {
+      atype,          n_node,
+      n_local,        canonicalIndexBytes(source),
+      edge_vec,       destination_row_ptr,
+      source_row_ptr, source_order};
+  if (long_range_) {
+    inputs.push_back(coord);
+    inputs.push_back(communicator_tensor(world));
+    inputs.push_back(cell);
+  }
+  const std::lock_guard evaluation(evaluation_mutex_);
+  return loader->run(inputs);
+}
+
+void DeepPotPTExpt::check_long_range_box(const std::vector& box) const {
+  if (!box.empty() && box.size() != 9) {
+    throw deepmd::deepmd_exception(
+        "a model with a long-range module evaluates one frame with a cell of "
+        "nine values, or without one, but got " +
+        std::to_string(box.size()) + " values.");
+  }
+  if (!box.empty() && !long_range_periodic_) {
+    throw deepmd::deepmd_exception(
+        "the long-range module of this model serves complete non-periodic "
+        "systems only; pass no cell (in LAMMPS, a non-periodic boundary f, s "
+        "or m in every dimension), or train the module with periodic: true.");
+  }
+  if (!box.empty()) {
+    // A cell whose volume is a vanishing fraction of the product of the
+    // lengths of its lattice vectors, or not finite, spans no lattice.
+    const double volume =
+        std::abs(box[0] * (box[4] * box[8] - box[5] * box[7]) -
+                 box[1] * (box[3] * box[8] - box[5] * box[6]) +
+                 box[2] * (box[3] * box[7] - box[4] * box[6]));
+    double lengths = 1.0;
+    for (int d = 0; d < 3; ++d) {
+      lengths *=
+          std::sqrt(box[3 * d] * box[3 * d] + box[3 * d + 1] * box[3 * d + 1] +
+                    box[3 * d + 2] * box[3 * d + 2]);
+    }
+    if (!(volume > 1e-6 * lengths)) {
+      throw deepmd::deepmd_exception(
+          "the cell of a periodic system needs three independent lattice "
+          "vectors, but its volume is " +
+          std::to_string(volume) + "; pass no cell for a non-periodic system.");
+    }
+  }
+}
+
+void DeepPotPTExpt::join_long_range_exchange(const void* world,
+                                             const std::vector& box) {
+  const torch::Device device(torch::kCUDA, gpu_id);
+  const auto i64 = torch::TensorOptions().dtype(torch::kInt64).device(device);
+  // The compact lower stores at least two edge slots; without edges both are
+  // padding outside every CSR row.
+  run_model_canonical_graph(
+      torch::zeros({1}, i64), torch::ones({1}, i64), torch::zeros({1}, i64),
+      torch::zeros({2}, i64.dtype(torch::kInt32)),
+      torch::zeros({2, 3}, i64.dtype(torch::kFloat32)), torch::zeros({2}, i64),
+      torch::zeros({2}, i64), torch::arange(2, i64),
+      torch::zeros({1, 3}, i64.dtype(torch::kFloat64)), world,
+      cell_tensor(box));
 }
 
 std::vector DeepPotPTExpt::run_model_with_comm(
@@ -799,6 +982,40 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener,
   bool multi_rank = (lmp_list.nprocs > 1);
   bool atom_map_present = (lmp_list.mapping != nullptr);
   bool use_with_comm = has_comm_artifact_ && multi_rank;
+  // A long-range module sums over every pair of atoms of a non-periodic
+  // system, which the caller marks with an empty box, and over the images of
+  // the cell of a periodic one. A single rank holds every atom, so a ghost
+  // atom there is a periodic image. Across ranks, the compact lower gathers
+  // the owned atoms of every rank through the communicator; the pairs that
+  // the graph lower lists stay within one.
+  const std::vector long_range_box(box.begin(), box.end());
+  if (long_range_) {
+    check_long_range_box(long_range_box);
+  }
+  if (long_range_ && long_range_box.empty() && !multi_rank && nghost_real > 0) {
+    throw deepmd::deepmd_exception(
+        "a model with a long-range module evaluates a complete non-periodic "
+        "system without a box, but this single-rank call carries " +
+        std::to_string(nghost_real) +
+        " ghost atoms, the periodic images of a periodic cell; pass the cell "
+        "of a periodic system to a model whose module serves one.");
+  }
+  if (long_range_ && multi_rank && !lower_input_is_canonical_) {
+    throw deepmd::deepmd_exception(
+        "a model with a long-range module evaluates a system decomposed over " +
+        std::to_string(lmp_list.nprocs) +
+        " MPI ranks only through the compact canonical lower, which a "
+        "compressed DPA4C model frozen for a CUDA device uses when the "
+        "long-range operator serves its long-range module; this archive uses "
+        "another lower, so run it on a single rank.");
+  }
+  if (long_range_ && multi_rank && lmp_list.world == nullptr) {
+    throw deepmd::deepmd_exception(
+        "a model with a long-range module evaluates a system decomposed over "
+        "several MPI ranks through their communicator, but the neighbor list "
+        "carries none (InputNlist::world).");
+  }
+  const void* const world = multi_rank ? lmp_list.world : nullptr;
   // Whether the edge topology folds ghost neighbours onto their local owners,
   // which reads an owner for every extended atom out of ``mapping_``.  The
   // edge lower folds unless the with-comm artifact carries ghost features
@@ -1227,7 +1444,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener,
         // ``Dim("n_node_total", min=1)``. Such a rank contributes nothing, so
         // fill zero outputs and return instead of running the model.  (The
         // tested ``nloc == 0`` empty-subdomain case has ``nall_real > 0`` --
-        // ghosts within rcut -- so it still runs the model normally.)
+        // ghosts within rcut -- so it still runs the model normally.) A
+        // long-range module across ranks still joins the exchanges of the
+        // others.
+        if (long_range_ && multi_rank) {
+          join_long_range_exchange(world, long_range_box);
+        }
         ener.assign(nframes, static_cast(0));
         force.assign(static_cast(nframes) * fwd_map.size() * 3,
                      static_cast(0));
@@ -1268,19 +1490,22 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener,
       graph_pack.atype = node_atype;
       graph_pack.n_node = n_node_tensor;
       graph_pack.n_local = n_local_tensor;
+      const at::Tensor node_coord =
+          coord_Tensor.slice(1, 0, n_node_count).reshape({n_node_count, 3});
       if (lower_input_is_canonical_) {
         const auto compact = compactCanonicalGraph(graph_pack);
         flat_outputs = run_model_canonical_graph(
             compact.atype, compact.n_node, compact.n_local, compact.source,
             compact.edge_vec, compact.destination_row_ptr,
-            compact.source_row_ptr, compact.source_order);
+            compact.source_row_ptr, compact.source_order, node_coord, world,
+            long_range_ ? cell_tensor(long_range_box) : torch::Tensor());
       } else {
         flat_outputs = run_model_graph(
             node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index,
             graph_pack.edge_vec, graph_pack.edge_mask,
             graph_pack.destination_order, graph_pack.destination_row_ptr,
             graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor,
-            graph_aparam, charge_spin_tensor);
+            graph_aparam, charge_spin_tensor, node_coord, long_range_box);
       }
     } else {
       // Model-level pair exclusion is a BUILD-time transform (decision
@@ -1473,6 +1698,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener,
   check_call_charge_spin(charge_spin, /*nframes=*/1, settable_chgspin,
                          dchgspin > 0, default_chg_spin_,
                          chg_spin_table_ranges_);
+  // A long-range module sums over the atoms of a non-periodic system, which
+  // the empty box marks, and over the images of the cell of a periodic one.
+  const std::vector long_range_box(box.begin(), box.end());
+  if (long_range_) {
+    check_long_range_box(long_range_box);
+  }
   // The .pt2 model only contains forward_common_lower, which requires
   // nlist as input. We must build the nlist in C++ and fold back the
   // extended-region outputs to local atoms.
@@ -1648,12 +1879,15 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener,
     if (graph_edge_fp32_) {
       graph_tensors.edge_vec = graph_tensors.edge_vec.to(torch::kFloat32);
     }
+    const at::Tensor node_coord =
+        coord_Tensor.slice(1, 0, nloc).reshape({nloc, 3});
     if (lower_input_is_canonical_) {
       const auto compact = compactCanonicalGraph(graph_tensors);
       flat_outputs = run_model_canonical_graph(
           compact.atype, compact.n_node, compact.n_local, compact.source,
           compact.edge_vec, compact.destination_row_ptr, compact.source_row_ptr,
-          compact.source_order);
+          compact.source_order, node_coord, /*world=*/nullptr,
+          long_range_ ? cell_tensor(long_range_box) : torch::Tensor());
     } else {
       flat_outputs = run_model_graph(
           graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local,
@@ -1665,7 +1899,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener,
           // normalizes the rectangular runtime aparam to the flat
           // (N, daparam) node axis of the graph ABI.
           extend_graph_aparam(aparam_tensor, natoms, natoms, daparam),
-          charge_spin_tensor);
+          charge_spin_tensor, node_coord, long_range_box);
     }
   } else {
     // Model-level pair exclusion is a BUILD-time transform (decision
@@ -2221,6 +2455,15 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy,
   const int nnode = (nall_nodes > 0) ? nall_nodes : nloc;
   const bool extended = nnode != nloc;
   const bool use_with_comm = has_comm_artifact_ && comm_nlist;
+  if (long_range_ && (extended || comm_nlist != nullptr)) {
+    throw deepmd::deepmd_exception(
+        "compute_edges_gpu: a model with a long-range module evaluates a "
+        "complete non-periodic system on a single rank (nall_nodes == nloc and "
+        "no comm_nlist); periodic systems and several ranks require the "
+        "compact canonical lower, which a compressed DPA4C model frozen for a "
+        "CUDA device uses when the long-range operator serves its module, or "
+        "the evaluation from a neighbor list (pair_style deepmd).");
+  }
   if (lower_input_is_graph_ && extended && has_message_passing_) {
     throw deepmd::deepmd_exception(
         "compute_edges_gpu: message-passing graph-input inference requires an "
@@ -2448,7 +2691,8 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy,
               graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask,
               graph_pack.destination_order, graph_pack.destination_row_ptr,
               graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor,
-              graph_aparam, charge_spin_tensor));
+              graph_aparam, charge_spin_tensor, coord_t.reshape({nloc, 3}),
+              /*box=*/{}));
       ae = out["atom_energy"].reshape({nnode}).slice(0, 0, nloc).contiguous();
       force_t = out["force"].reshape({nnode, 3}).contiguous();
       av = out["atom_virial"].reshape({nnode, 9}).contiguous();
@@ -2528,23 +2772,29 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy,
       torch::from_blob(d_force, {nnode, 3}, opt_f64).copy_(force_t);
       torch::from_blob(d_atom_virial, {nnode, 9}, opt_f64).copy_(av);
     }
-    synchronize_current_accelerator_stream();
+    synchronize_current_accelerator_stream(gpu_id);
   });
 }
 
 void DeepPotPTExpt::compute_canonical_graph_gpu_impl(
     double* d_atom_energy,
     double* d_force,
+    double* d_virial,
     double* d_atom_virial,
     const std::int64_t* d_atype,
-    const std::uint32_t* d_source,
+    const void* d_source,
+    DP_CanonicalIndexWidth index_width,
     const float* d_edge_vec,
     const std::int64_t* d_destination_row_ptr,
     const std::int64_t* d_source_row_ptr,
-    const std::uint32_t* d_source_order,
-    const int nloc,
-    const int nall_nodes,
-    const std::int64_t edge_storage) {
+    const std::int64_t* d_source_order,
+    const double* d_coord,
+    void* world,
+    const double* box,
+    const std::int64_t nloc,
+    const std::int64_t nall_nodes,
+    const std::int64_t edge_storage,
+    void* stream) {
   if (!lower_input_is_canonical_) {
     throw deepmd::deepmd_exception(
         "compute_canonical_graph_gpu requires a compact canonical artifact.");
@@ -2553,66 +2803,236 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl(
     throw deepmd::deepmd_exception(
         "compute_canonical_graph_gpu requires a CUDA device.");
   }
-  if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2 ||
-      static_cast(edge_storage) >
-          std::numeric_limits::max()) {
+  // A long-range module reads the positions of a complete system and the
+  // cell of a periodic one: the node set of a single rank, whose halo nodes
+  // are images of a periodic cell, or the owned nodes of every rank of
+  // ``world``, which a rank without nodes still joins.
+  const std::vector long_range_box =
+      box != nullptr ? std::vector(box, box + 9)
+                     : std::vector();
+  if (long_range_) {
+    check_long_range_box(long_range_box);
+  }
+  if (long_range_ && d_coord == nullptr && nall_nodes > 0) {
+    throw deepmd::deepmd_exception(
+        "compute_canonical_graph_gpu: a model with a long-range module reads "
+        "the node positions (d_coord).");
+  }
+  if (long_range_ && world == nullptr && box == nullptr && nall_nodes != nloc) {
+    throw deepmd::deepmd_exception(
+        "compute_canonical_graph_gpu: a model with a long-range module on a "
+        "single rank evaluates a complete non-periodic system "
+        "(nall_nodes == nloc), or a periodic one with its cell; a system "
+        "decomposed over several ranks takes their communicator.");
+  }
+  if (long_range_ && world != nullptr && nall_nodes == 0) {
+    // A rank without nodes owns none and stores no edge.
+    if (nloc != 0 || edge_storage != 0) {
+      throw deepmd::deepmd_exception(
+          "invalid compact canonical graph dimensions.");
+    }
+    translate_error([&] {
+      const c10::DeviceGuard device_guard(torch::Device(torch::kCUDA, gpu_id));
+      const ptexpt::CallerStreamScope stream_scope(stream, gpu_id);
+      join_long_range_exchange(world, long_range_box);
+      if (stream == nullptr) {
+        synchronize_current_accelerator_stream(gpu_id);
+      }
+    });
+    return;
+  }
+  if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2) {
     throw deepmd::deepmd_exception(
         "invalid compact canonical graph dimensions.");
   }
+  ptexpt::check_canonical_index(d_source, index_width, nall_nodes);
+  if (d_force == nullptr || d_atype == nullptr || d_source == nullptr ||
+      d_edge_vec == nullptr || d_destination_row_ptr == nullptr ||
+      (nloc > 0 && d_atom_energy == nullptr)) {
+    throw deepmd::deepmd_exception(
+        "compute_canonical_graph_gpu: the energy and force outputs and the "
+        "graph inputs must not be null; only the virial, the per-atom virial "
+        "and the source CSR are optional.");
+  }
+  if (d_atom_virial != nullptr && !do_atomic_virial) {
+    throw deepmd::deepmd_exception(
+        "compute_canonical_graph_gpu: the per-atom virial was requested, but "
+        "this .pt2 model was exported without it (do_atomic_virial=False).");
+  }
+  const bool has_source_csr =
+      d_source_row_ptr != nullptr && d_source_order != nullptr;
+  if (!has_source_csr &&
+      (d_source_row_ptr != nullptr || d_source_order != nullptr ||
+       d_atom_virial != nullptr || canonical_requires_source_csr_)) {
+    throw deepmd::deepmd_exception(
+        "compute_canonical_graph_gpu: the source CSR may be omitted only as "
+        "a pair, without a per-atom virial, and only for an artifact that "
+        "assembles the force without it.");
+  }
 
   translate_error([&] {
+    // The source CSR placeholders and the replay belong to the instance, so
+    // one evaluation at a time prepares and runs them.
+    const std::lock_guard evaluation(evaluation_mutex_);
     const torch::Device device(torch::kCUDA, gpu_id);
     const c10::DeviceGuard device_guard(device);
-    const auto opt_f32 =
-        torch::TensorOptions().dtype(torch::kFloat32).device(device);
+    const ptexpt::CallerStreamScope stream_scope(stream, gpu_id);
     const auto opt_i64 =
         torch::TensorOptions().dtype(torch::kInt64).device(device);
-    const auto opt_u32 = torch::TensorOptions()
-                             .dtype(deepmd::canonicalGraphIndexType())
-                             .device(device);
-    auto atype = torch::from_blob(const_cast(d_atype),
-                                  {nall_nodes}, opt_i64);
-    auto source = torch::from_blob(const_cast(d_source),
-                                   {edge_storage}, opt_u32);
-    auto edge_vec = torch::from_blob(const_cast(d_edge_vec),
-                                     {edge_storage, 3}, opt_f32);
-    auto destination_row_ptr =
-        torch::from_blob(const_cast(d_destination_row_ptr),
-                         {nall_nodes + 1}, opt_i64);
-    auto source_row_ptr = torch::from_blob(
-        const_cast(d_source_row_ptr), {nall_nodes + 1}, opt_i64);
-    auto source_order = torch::from_blob(
-        const_cast(d_source_order), {edge_storage}, opt_u32);
-    auto n_node = torch::full({1}, nall_nodes, opt_i64);
-    auto n_local = torch::full({1}, nloc, opt_i64);
-
-    std::map output;
-    extract_outputs(output,
-                    run_model_canonical_graph(atype, n_node, n_local, source,
-                                              edge_vec, destination_row_ptr,
-                                              source_row_ptr, source_order));
-    auto atom_energy = output["atom_energy"]
-                           .reshape({nall_nodes})
-                           .slice(0, 0, nloc)
-                           .contiguous();
-    auto force = output["force"].reshape({nall_nodes, 3}).contiguous();
-    auto atom_virial =
-        output["atom_virial"].reshape({nall_nodes, 9}).contiguous();
-    if (nloc > 0) {
-      torch::from_blob(
-          d_atom_energy, {nloc},
-          torch::TensorOptions().dtype(torch::kFloat64).device(device))
-          .copy_(atom_energy);
+    const auto opt_index =
+        torch::TensorOptions()
+            .dtype(index_width == DP_CANONICAL_INDEX32 ? torch::kInt32
+                                                       : torch::kInt64)
+            .device(device);
+    if (!has_source_csr && !source_row_placeholder_.defined()) {
+      source_row_placeholder_ = torch::zeros({2}, opt_i64);
+      source_order_placeholder_ = torch::zeros({2}, opt_i64);
+    }
+    const void* const source_row_data =
+        has_source_csr ? d_source_row_ptr : source_row_placeholder_.data_ptr();
+    const void* const source_order_data =
+        has_source_csr ? d_source_order : source_order_placeholder_.data_ptr();
+    // A model with a long-range module takes the cell in pinned memory, which
+    // the evaluation of a periodic system reads on the device.
+    const bool long_range = has_long_range();
+    if (long_range && !long_range_cell_) {
+      long_range_cell_ = std::make_unique();
+    }
+    const torch::Tensor cell =
+        long_range ? long_range_cell_->hold(long_range_box, gpu_id)
+                   : torch::Tensor();
+
+    // The tensors over the caller's buffers are built inside the evaluation,
+    // which a replay skips: wrapping a device buffer queries the driver for
+    // the device that holds it.
+    const auto evaluate = [&] {
+      const auto opt_f32 =
+          torch::TensorOptions().dtype(torch::kFloat32).device(device);
+      const auto opt_f64 =
+          torch::TensorOptions().dtype(torch::kFloat64).device(device);
+      auto atype = torch::from_blob(const_cast(d_atype),
+                                    {nall_nodes}, opt_i64);
+      auto source = torch::from_blob(const_cast(d_source),
+                                     {edge_storage}, opt_index);
+      auto edge_vec = torch::from_blob(const_cast(d_edge_vec),
+                                       {edge_storage, 3}, opt_f32);
+      auto destination_row_ptr =
+          torch::from_blob(const_cast(d_destination_row_ptr),
+                           {nall_nodes + 1}, opt_i64);
+      auto source_row_ptr =
+          has_source_csr
+              ? torch::from_blob(const_cast(d_source_row_ptr),
+                                 {nall_nodes + 1}, opt_i64)
+              : source_row_placeholder_;
+      auto source_order =
+          has_source_csr
+              ? torch::from_blob(const_cast(d_source_order),
+                                 {edge_storage}, opt_i64)
+              : source_order_placeholder_;
+      // The node counts are allocated inside the evaluation, so that a
+      // captured graph owns them.
+      auto n_node = torch::full({1}, nall_nodes, opt_i64);
+      auto n_local = torch::full({1}, nloc, opt_i64);
+      const auto coord = d_coord != nullptr
+                             ? torch::from_blob(const_cast(d_coord),
+                                                {nall_nodes, 3}, opt_f64)
+                             : torch::Tensor();
+      std::map output;
+      extract_outputs(output, run_model_canonical_graph(
+                                  atype, n_node, n_local, source, edge_vec,
+                                  destination_row_ptr, source_row_ptr,
+                                  source_order, coord, world, cell));
+      auto atom_energy = output["atom_energy"]
+                             .reshape({nall_nodes})
+                             .slice(0, 0, nloc)
+                             .contiguous();
+      auto force = output["force"].reshape({nall_nodes, 3}).contiguous();
+      // The results reach the fp64 outputs through ptexpt::copy_results. The
+      // DPA4C operator returns the force, the virial and the atomic energy in
+      // fp64, a list of one dtype; the DPA1 assembly returns its force and
+      // virial in fp32 beside an fp64 atomic energy, which the copy converts
+      // one tensor at a time.
+      std::vector targets = {
+          torch::from_blob(d_force, {nall_nodes, 3}, opt_f64)};
+      std::vector results = {force};
+      if (d_virial != nullptr) {
+        targets.push_back(torch::from_blob(d_virial, {9}, opt_f64));
+        results.push_back(output["virial"].reshape({9}));
+      }
+      if (nloc > 0) {
+        targets.push_back(torch::from_blob(d_atom_energy, {nloc}, opt_f64));
+        results.push_back(atom_energy);
+      }
+      ptexpt::copy_results(targets, results, nall_nodes);
+      if (d_atom_virial != nullptr) {
+        torch::from_blob(d_atom_virial, {nall_nodes, 9}, opt_f64)
+            .copy_(output["atom_virial"].reshape({nall_nodes, 9}));
+      }
+    };
+    // A graph without halo nodes keeps its extents from step to step, so its
+    // evaluation is replayed; a domain's halo changes its extent at every
+    // neighbor rebuild, so a graph with halo nodes is evaluated eagerly. A
+    // model with a long-range module reads the node positions, whose address
+    // is then part of the key, and is evaluated eagerly over several ranks,
+    // where it may communicate between its kernels. A replay reads the cell
+    // of its call, while the mesh size of the cell fixes the buffers and the
+    // launches of the evaluation, so the size joins the key.
+    bool replay = nloc == nall_nodes && !(long_range && world != nullptr);
+    std::array mesh = {};
+    if (long_range) {
+      if (!long_range_box.empty()) {
+        mesh = ptexpt::meshDims(long_range_box, mesh_spacing_);
+      }
+      if (long_range_mesh_ != mesh) {
+        // The first mesh of the instance is replayed once its key repeats. The
+        // graph of another mesh is released at once, so that its memory is
+        // free during the eager calls that follow.
+        long_range_mesh_calls_ =
+            long_range_mesh_.has_value() ? 0 : kStableMeshCalls;
+        long_range_mesh_ = mesh;
+        if (canonical_replay_) {
+          canonical_replay_->invalidate();
+        }
+      }
+      if (long_range_mesh_calls_ < kStableMeshCalls) {
+        ++long_range_mesh_calls_;
+        replay = false;
+      }
+    }
+    if (!replay) {
+      evaluate();
+    } else {
+      if (!canonical_replay_) {
+        canonical_replay_ = std::make_unique();
+      }
+      ptexpt::GraphReplay::Key key = {
+          ptexpt::address_key(d_atom_energy),
+          ptexpt::address_key(d_force),
+          ptexpt::address_key(d_virial),
+          ptexpt::address_key(d_atom_virial),
+          ptexpt::address_key(d_atype),
+          ptexpt::address_key(d_source),
+          ptexpt::address_key(d_edge_vec),
+          ptexpt::address_key(d_destination_row_ptr),
+          ptexpt::address_key(source_row_data),
+          ptexpt::address_key(source_order_data),
+          ptexpt::address_key(long_range ? d_coord : nullptr),
+          nloc,
+          nall_nodes,
+          edge_storage,
+          index_width};
+      if (long_range) {
+        key.push_back(ptexpt::address_key(cell.data_ptr()));
+        key.insert(key.end(), mesh.begin(), mesh.end());
+      }
+      canonical_replay_->run(key, gpu_id, evaluate);
+    }
+    if (long_range_periodic_) {
+      long_range_cell_->record();
+    }
+    if (stream == nullptr) {
+      synchronize_current_accelerator_stream(gpu_id);
     }
-    torch::from_blob(
-        d_force, {nall_nodes, 3},
-        torch::TensorOptions().dtype(torch::kFloat64).device(device))
-        .copy_(force);
-    torch::from_blob(
-        d_atom_virial, {nall_nodes, 9},
-        torch::TensorOptions().dtype(torch::kFloat64).device(device))
-        .copy_(atom_virial);
-    synchronize_current_accelerator_stream();
   });
 }
 
@@ -2670,20 +3090,27 @@ void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy,
 void DeepPotPTExpt::compute_canonical_graph_gpu(
     double* d_atom_energy,
     double* d_force,
+    double* d_virial,
     double* d_atom_virial,
     const std::int64_t* d_atype,
-    const std::uint32_t* d_source,
+    const void* d_source,
+    DP_CanonicalIndexWidth index_width,
     const float* d_edge_vec,
     const std::int64_t* d_destination_row_ptr,
     const std::int64_t* d_source_row_ptr,
-    const std::uint32_t* d_source_order,
-    const int nloc,
-    const int nall_nodes,
-    const std::int64_t edge_storage) {
+    const std::int64_t* d_source_order,
+    const double* d_coord,
+    void* world,
+    const double* box,
+    const std::int64_t nloc,
+    const std::int64_t nall_nodes,
+    const std::int64_t edge_storage,
+    void* stream) {
   compute_canonical_graph_gpu_impl(
-      d_atom_energy, d_force, d_atom_virial, d_atype, d_source, d_edge_vec,
-      d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes,
-      edge_storage);
+      d_atom_energy, d_force, d_virial, d_atom_virial, d_atype, d_source,
+      index_width, d_edge_vec, d_destination_row_ptr, d_source_row_ptr,
+      d_source_order, d_coord, world, box, nloc, nall_nodes, edge_storage,
+      stream);
 }
 
 bool DeepPotPTExpt::uses_fp32_edge_vectors() const {
@@ -2700,4 +3127,10 @@ bool DeepPotPTExpt::uses_canonical_graph_inference() const {
   return lower_input_is_canonical_;
 }
 
+bool DeepPotPTExpt::canonical_graph_requires_source_csr() const {
+  return canonical_requires_source_csr_;
+}
+
+bool DeepPotPTExpt::has_long_range() const { return long_range_; }
+
 #endif
diff --git a/source/api_cc/src/DeepPotPTExptPlugin.cc b/source/api_cc/src/DeepPotPTExptPlugin.cc
index 3499c62dcb..da6b6f93b0 100644
--- a/source/api_cc/src/DeepPotPTExptPlugin.cc
+++ b/source/api_cc/src/DeepPotPTExptPlugin.cc
@@ -57,7 +57,7 @@ bool native_spin_backend_serves(const char* model) {
 }  // namespace
 #endif
 
-extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
+extern "C" void* deepmd_create_deeppot_backend_v2(const char* model,
                                                   int gpu_rank,
                                                   const char* file_content,
                                                   std::size_t file_content_size,
@@ -74,11 +74,11 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
 #endif
 }
 
-extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) {
+extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) {
   deepmd::plugin::delete_deeppot_backend(backend);
 }
 
-extern "C" void* deepmd_create_deepspin_backend_v1(
+extern "C" void* deepmd_create_deepspin_backend_v2(
     const char* model,
     int gpu_rank,
     const char* file_content,
@@ -102,7 +102,7 @@ extern "C" void* deepmd_create_deepspin_backend_v1(
 #endif
 }
 
-extern "C" void deepmd_delete_deepspin_backend_v1(void* backend) {
+extern "C" void deepmd_delete_deepspin_backend_v2(void* backend) {
   deepmd::plugin::delete_deepspin_backend(backend);
 }
 
diff --git a/source/api_cc/src/DeepPotPTPlugin.cc b/source/api_cc/src/DeepPotPTPlugin.cc
index ae584b12b6..f024c4e10a 100644
--- a/source/api_cc/src/DeepPotPTPlugin.cc
+++ b/source/api_cc/src/DeepPotPTPlugin.cc
@@ -6,7 +6,7 @@
 #include "DeepSpinPT.h"
 #include "DeepTensorPT.h"
 
-extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
+extern "C" void* deepmd_create_deeppot_backend_v2(const char* model,
                                                   int gpu_rank,
                                                   const char* file_content,
                                                   std::size_t file_content_size,
@@ -15,11 +15,11 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
       model, gpu_rank, file_content, file_content_size, error_message);
 }
 
-extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) {
+extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) {
   deepmd::plugin::delete_deeppot_backend(backend);
 }
 
-extern "C" void* deepmd_create_deepspin_backend_v1(
+extern "C" void* deepmd_create_deepspin_backend_v2(
     const char* model,
     int gpu_rank,
     const char* file_content,
@@ -29,7 +29,7 @@ extern "C" void* deepmd_create_deepspin_backend_v1(
       model, gpu_rank, file_content, file_content_size, error_message);
 }
 
-extern "C" void deepmd_delete_deepspin_backend_v1(void* backend) {
+extern "C" void deepmd_delete_deepspin_backend_v2(void* backend) {
   deepmd::plugin::delete_deepspin_backend(backend);
 }
 
diff --git a/source/api_cc/src/DeepPotTFPlugin.cc b/source/api_cc/src/DeepPotTFPlugin.cc
index 5e20b6fcf7..f1daf082a9 100644
--- a/source/api_cc/src/DeepPotTFPlugin.cc
+++ b/source/api_cc/src/DeepPotTFPlugin.cc
@@ -46,7 +46,7 @@ void convert_pbtxt_to_pb_impl(const char* pbtxt, const char* pb) {
 
 }  // namespace
 
-extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
+extern "C" void* deepmd_create_deeppot_backend_v2(const char* model,
                                                   int gpu_rank,
                                                   const char* file_content,
                                                   std::size_t file_content_size,
@@ -55,11 +55,11 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model,
       model, gpu_rank, file_content, file_content_size, error_message);
 }
 
-extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) {
+extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) {
   deepmd::plugin::delete_deeppot_backend(backend);
 }
 
-extern "C" void* deepmd_create_deepspin_backend_v1(
+extern "C" void* deepmd_create_deepspin_backend_v2(
     const char* model,
     int gpu_rank,
     const char* file_content,
@@ -69,7 +69,7 @@ extern "C" void* deepmd_create_deepspin_backend_v1(
       model, gpu_rank, file_content, file_content_size, error_message);
 }
 
-extern "C" void deepmd_delete_deepspin_backend_v1(void* backend) {
+extern "C" void deepmd_delete_deepspin_backend_v2(void* backend) {
   deepmd::plugin::delete_deepspin_backend(backend);
 }
 
diff --git a/source/api_cc/src/DeepSpin.cc b/source/api_cc/src/DeepSpin.cc
index e52d6a657b..83aa44c08b 100644
--- a/source/api_cc/src/DeepSpin.cc
+++ b/source/api_cc/src/DeepSpin.cc
@@ -473,20 +473,24 @@ void DeepSpinBackend::compute_canonical_graph_gpu(
     double* d_atom_energy,
     double* d_force,
     double* d_force_mag,
+    double* d_virial,
     double* d_atom_virial,
     const std::int64_t* d_atype,
-    const std::uint32_t* d_source,
+    const void* d_source,
+    DP_CanonicalIndexWidth index_width,
     const float* d_edge_vec,
     const std::int64_t* d_destination_row_ptr,
     const std::int64_t* d_source_row_ptr,
-    const std::uint32_t* d_source_order,
+    const std::int64_t* d_source_order,
     const float* d_spin,
-    const int nloc,
-    const int nall_nodes,
-    const std::int64_t edge_storage) {
+    const std::int64_t nloc,
+    const std::int64_t nall_nodes,
+    const std::int64_t edge_storage,
+    void* stream) {
   (void)d_atom_energy;
   (void)d_force;
   (void)d_force_mag;
+  (void)d_virial;
   (void)d_atom_virial;
   (void)d_atype;
   (void)d_source;
@@ -498,6 +502,7 @@ void DeepSpinBackend::compute_canonical_graph_gpu(
   (void)nloc;
   (void)nall_nodes;
   (void)edge_storage;
+  (void)stream;
   throw deepmd::deepmd_exception(
       "compact canonical graph inference is only supported by a compatible "
       "PyTorch Exportable backend.");
@@ -511,25 +516,29 @@ void DeepSpin::compute_canonical_graph_gpu(
     double* d_atom_energy,
     double* d_force,
     double* d_force_mag,
+    double* d_virial,
     double* d_atom_virial,
     const std::int64_t* d_atype,
-    const std::uint32_t* d_source,
+    const void* d_source,
+    DP_CanonicalIndexWidth index_width,
     const float* d_edge_vec,
     const std::int64_t* d_destination_row_ptr,
     const std::int64_t* d_source_row_ptr,
-    const std::uint32_t* d_source_order,
+    const std::int64_t* d_source_order,
     const float* d_spin,
-    const int nloc,
-    const int nall_nodes,
-    const std::int64_t edge_storage) {
+    const std::int64_t nloc,
+    const std::int64_t nall_nodes,
+    const std::int64_t edge_storage,
+    void* stream) {
   // Backend-agnostic dispatch: backends that implement device-resident
   // canonical inference override the hook, while the others inherit the
   // throwing default. ``libdeepmd_cc`` does not link any backend, so the
   // dispatch stays virtual rather than casting to a concrete backend type.
-  dp->compute_canonical_graph_gpu(
-      d_atom_energy, d_force, d_force_mag, d_atom_virial, d_atype, d_source,
-      d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order,
-      d_spin, nloc, nall_nodes, edge_storage);
+  dp->compute_canonical_graph_gpu(d_atom_energy, d_force, d_force_mag, d_virial,
+                                  d_atom_virial, d_atype, d_source, index_width,
+                                  d_edge_vec, d_destination_row_ptr,
+                                  d_source_row_ptr, d_source_order, d_spin,
+                                  nloc, nall_nodes, edge_storage, stream);
 }
 
 bool DeepSpin::uses_canonical_graph_inference() const {
diff --git a/source/api_cc/src/NativeSpinPTExpt.cc b/source/api_cc/src/NativeSpinPTExpt.cc
index ec684b26f8..14ba150421 100644
--- a/source/api_cc/src/NativeSpinPTExpt.cc
+++ b/source/api_cc/src/NativeSpinPTExpt.cc
@@ -32,10 +32,19 @@ using namespace deepmd;
 
 namespace {
 
-void synchronize_current_accelerator_stream() {
-#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
+// Wait for the evaluation that the calling thread enqueued on its current
+// stream of ``device``, on which the device-resident entry points without a
+// caller's stream enqueue all of their work; the work and the CUDA graph
+// captures of other threads on the device proceed meanwhile.
+void synchronize_current_accelerator_stream(const int device) {
+#if defined(DEEPMD_TORCH_HAS_CUDA)
+  c10::cuda::getCurrentCUDAStream(static_cast(device))
+      .synchronize();
+#elif defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
+  (void)device;
   DPErrcheck(gpuDeviceSynchronize());
 #else
+  (void)device;
   throw deepmd::deepmd_exception(
       "GPU-resident inference requires a GPU-enabled DeePMD-kit build.");
 #endif
@@ -381,9 +390,12 @@ void NativeSpinPTExpt::init(const std::string& model,
           "compact canonical graph artifacts require float32 edge vectors.");
     }
     if (!metadata.obj_val.count("canonical_index_dtype") ||
-        metadata["canonical_index_dtype"].as_string() != "uint32") {
+        metadata["canonical_index_dtype"].as_string() != "auto" ||
+        !metadata.obj_val.count("canonical_index_storage") ||
+        metadata["canonical_index_storage"].as_string() != "byte_words") {
       throw deepmd::deepmd_exception(
-          "compact canonical graph artifacts require uint32 topology; "
+          "compact canonical graph artifacts require automatic-width byte "
+          "indices; "
           "re-freeze the model with the current DeePMD-kit version.");
     }
     // The compact lower is traced with the nine graph and moment inputs
@@ -399,9 +411,9 @@ void NativeSpinPTExpt::init(const std::string& model,
           "; freeze it with the graph lower instead.");
     }
   }
-  // The per-atom virial is a structural part of both contracts: the
-  // device-resident entry point returns it unconditionally and the host paths
-  // reduce it into the global virial.
+  // The per-atom virial is a structural part of the host contract, whose paths
+  // reduce it into the global virial; the device-resident entry point returns
+  // it whenever it is requested.
   if (!metadata.obj_val.count("do_atomic_virial") ||
       !metadata["do_atomic_virial"].as_bool()) {
     throw deepmd::deepmd_exception(
@@ -436,8 +448,12 @@ void NativeSpinPTExpt::init(const std::string& model,
     output_keys.push_back(v.as_string());
   }
 
+  // An instance runs its model one call at a time (evaluation_mutex_), so the
+  // model runs on the single-threaded runner, which records no completion
+  // event per run: a run is free to be enqueued before the previous one
+  // finishes and to be captured into a CUDA graph.
   loader = std::make_unique(
-      model, "model", false, 1,
+      model, "model", /*run_single_threaded=*/true, 1,
       gpu_enabled ? static_cast(gpu_id)
                   : static_cast(-1));
 
@@ -530,10 +546,27 @@ void NativeSpinPTExpt::set_charge_spin(const std::vector& charge_spin) {
   // constants without opening a second failure point.
   std::vector next_charge_spin(charge_spin);
   if (charge_state_fold_) {
-    charge_state_fold_->apply(charge_spin,
-                              gpu_enabled ? torch::Device(torch::kCUDA, gpu_id)
-                                          : torch::Device(torch::kCPU),
-                              *loader);
+    const std::lock_guard evaluation(evaluation_mutex_);
+    translate_error([&] {
+      // Evaluations still in flight on a caller's stream, and a captured
+      // evaluation, read the constants the fold replaces, and the next
+      // evaluation, possibly on another stream, reads those it writes; the
+      // fold therefore waits for the device on both sides.
+      if (gpu_enabled) {
+        ptexpt::synchronize_device(gpu_id);
+      }
+      if (canonical_replay_) {
+        canonical_replay_->invalidate();
+      }
+      charge_state_fold_->apply(charge_spin,
+                                gpu_enabled
+                                    ? torch::Device(torch::kCUDA, gpu_id)
+                                    : torch::Device(torch::kCPU),
+                                *loader);
+      if (gpu_enabled) {
+        ptexpt::synchronize_device(gpu_id);
+      }
+    });
   }
   default_chg_spin_.swap(next_charge_spin);
 }
@@ -552,10 +585,11 @@ std::vector NativeSpinPTExpt::run_model_canonical(
     const CanonicalGraphTensorPack& graph, const torch::Tensor& spin) {
   // The moment shares the float32 precision of the compact geometry; the cast
   // is a no-op for a caller that already holds a float32 tensor.
-  return loader->run({graph.atype, graph.n_node, graph.n_local, graph.source,
-                      graph.edge_vec, graph.destination_row_ptr,
-                      graph.source_row_ptr, graph.source_order,
-                      spin.to(torch::kFloat32)});
+  const std::lock_guard evaluation(evaluation_mutex_);
+  return loader->run({graph.atype, graph.n_node, graph.n_local,
+                      canonicalIndexBytes(graph.source), graph.edge_vec,
+                      graph.destination_row_ptr, graph.source_row_ptr,
+                      graph.source_order, spin.to(torch::kFloat32)});
 }
 
 std::vector NativeSpinPTExpt::run_model_graph(
@@ -587,6 +621,7 @@ std::vector NativeSpinPTExpt::run_model_graph(
   if (dchgspin > 0) {
     inputs.push_back(charge_spin);
   }
+  const std::lock_guard evaluation(evaluation_mutex_);
   return loader->run(inputs);
 }
 
@@ -1333,17 +1368,20 @@ void NativeSpinPTExpt::compute_canonical_graph_gpu(
     double* d_atom_energy,
     double* d_force,
     double* d_force_mag,
+    double* d_virial,
     double* d_atom_virial,
     const std::int64_t* d_atype,
-    const std::uint32_t* d_source,
+    const void* d_source,
+    DP_CanonicalIndexWidth index_width,
     const float* d_edge_vec,
     const std::int64_t* d_destination_row_ptr,
     const std::int64_t* d_source_row_ptr,
-    const std::uint32_t* d_source_order,
+    const std::int64_t* d_source_order,
     const float* d_spin,
-    const int nloc,
-    const int nall_nodes,
-    const std::int64_t edge_storage) {
+    const std::int64_t nloc,
+    const std::int64_t nall_nodes,
+    const std::int64_t edge_storage,
+    void* stream) {
   if (!canonical_abi_) {
     throw deepmd::deepmd_exception(
         "device-resident inference consumes the compact canonical graph; this "
@@ -1354,62 +1392,140 @@ void NativeSpinPTExpt::compute_canonical_graph_gpu(
     throw deepmd::deepmd_exception(
         "compute_canonical_graph_gpu requires a CUDA device.");
   }
-  if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2 ||
-      static_cast(edge_storage) >
-          std::numeric_limits::max()) {
+  if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2) {
     throw deepmd::deepmd_exception(
         "invalid compact canonical graph dimensions.");
   }
+  ptexpt::check_canonical_index(d_source, index_width, nall_nodes);
+  if (d_force == nullptr || d_force_mag == nullptr || d_atype == nullptr ||
+      d_source == nullptr || d_edge_vec == nullptr ||
+      d_destination_row_ptr == nullptr || d_spin == nullptr ||
+      (nloc > 0 && d_atom_energy == nullptr)) {
+    throw deepmd::deepmd_exception(
+        "compute_canonical_graph_gpu: the energy, force and magnetic force "
+        "outputs and the graph and moment inputs must not be null; only the "
+        "virial, the per-atom virial and the source CSR are optional.");
+  }
+  const bool has_source_csr =
+      d_source_row_ptr != nullptr && d_source_order != nullptr;
+  if (!has_source_csr &&
+      (d_source_row_ptr != nullptr || d_source_order != nullptr ||
+       d_atom_virial != nullptr)) {
+    throw deepmd::deepmd_exception(
+        "compute_canonical_graph_gpu: the source CSR may be omitted only as "
+        "a pair and without a per-atom virial.");
+  }
 
   translate_error([&] {
+    // The source CSR placeholders and the replay belong to the instance, so
+    // one evaluation at a time prepares and runs them.
+    const std::lock_guard evaluation(evaluation_mutex_);
     const torch::Device device(torch::kCUDA, gpu_id);
     const c10::DeviceGuard device_guard(device);
-    const auto opt_f32 =
-        torch::TensorOptions().dtype(torch::kFloat32).device(device);
-    const auto opt_f64 =
-        torch::TensorOptions().dtype(torch::kFloat64).device(device);
+    const ptexpt::CallerStreamScope stream_scope(stream, gpu_id);
     const auto opt_i64 =
         torch::TensorOptions().dtype(torch::kInt64).device(device);
-    const auto opt_u32 = torch::TensorOptions()
-                             .dtype(deepmd::canonicalGraphIndexType())
-                             .device(device);
-    CanonicalGraphTensorPack graph;
-    graph.atype = torch::from_blob(const_cast(d_atype),
-                                   {nall_nodes}, opt_i64);
-    graph.n_node = torch::full({1}, nall_nodes, opt_i64);
-    graph.n_local = torch::full({1}, nloc, opt_i64);
-    graph.source = torch::from_blob(const_cast(d_source),
-                                    {edge_storage}, opt_u32);
-    graph.edge_vec = torch::from_blob(const_cast(d_edge_vec),
-                                      {edge_storage, 3}, opt_f32);
-    graph.destination_row_ptr =
-        torch::from_blob(const_cast(d_destination_row_ptr),
-                         {nall_nodes + 1}, opt_i64);
-    graph.source_row_ptr = torch::from_blob(
-        const_cast(d_source_row_ptr), {nall_nodes + 1}, opt_i64);
-    graph.source_order = torch::from_blob(
-        const_cast(d_source_order), {edge_storage}, opt_u32);
-    const auto spin =
-        torch::from_blob(const_cast(d_spin), {nall_nodes, 3}, opt_f32);
-
-    std::map output;
-    extract_outputs(output, run_model_canonical(graph, spin));
-    auto atom_energy = output["atom_energy"]
-                           .reshape({nall_nodes})
-                           .slice(0, 0, nloc)
-                           .contiguous();
-    auto force = output["force"].reshape({nall_nodes, 3}).contiguous();
-    auto force_mag = output["force_mag"].reshape({nall_nodes, 3}).contiguous();
-    auto atom_virial =
-        output["atom_virial"].reshape({nall_nodes, 9}).contiguous();
-    if (nloc > 0) {
-      torch::from_blob(d_atom_energy, {nloc}, opt_f64).copy_(atom_energy);
+    const auto opt_index =
+        torch::TensorOptions()
+            .dtype(index_width == DP_CANONICAL_INDEX32 ? torch::kInt32
+                                                       : torch::kInt64)
+            .device(device);
+    if (!has_source_csr && !source_row_placeholder_.defined()) {
+      source_row_placeholder_ = torch::zeros({2}, opt_i64);
+      source_order_placeholder_ = torch::zeros({2}, opt_i64);
+    }
+    const void* const source_row_data =
+        has_source_csr ? d_source_row_ptr : source_row_placeholder_.data_ptr();
+    const void* const source_order_data =
+        has_source_csr ? d_source_order : source_order_placeholder_.data_ptr();
+
+    // The tensors over the caller's buffers are built inside the evaluation,
+    // which a replay skips: wrapping a device buffer queries the driver for
+    // the device that holds it.
+    const auto evaluate = [&] {
+      const auto opt_f32 =
+          torch::TensorOptions().dtype(torch::kFloat32).device(device);
+      const auto opt_f64 =
+          torch::TensorOptions().dtype(torch::kFloat64).device(device);
+      CanonicalGraphTensorPack inputs;
+      inputs.atype = torch::from_blob(const_cast(d_atype),
+                                      {nall_nodes}, opt_i64);
+      inputs.source = torch::from_blob(const_cast(d_source),
+                                       {edge_storage}, opt_index);
+      inputs.edge_vec = torch::from_blob(const_cast(d_edge_vec),
+                                         {edge_storage, 3}, opt_f32);
+      inputs.destination_row_ptr =
+          torch::from_blob(const_cast(d_destination_row_ptr),
+                           {nall_nodes + 1}, opt_i64);
+      inputs.source_row_ptr =
+          has_source_csr
+              ? torch::from_blob(const_cast(d_source_row_ptr),
+                                 {nall_nodes + 1}, opt_i64)
+              : source_row_placeholder_;
+      inputs.source_order =
+          has_source_csr
+              ? torch::from_blob(const_cast(d_source_order),
+                                 {edge_storage}, opt_i64)
+              : source_order_placeholder_;
+      const auto spin = torch::from_blob(const_cast(d_spin),
+                                         {nall_nodes, 3}, opt_f32);
+      // The node counts are allocated inside the evaluation, so that a
+      // captured graph owns them.
+      inputs.n_node = torch::full({1}, nall_nodes, opt_i64);
+      inputs.n_local = torch::full({1}, nloc, opt_i64);
+      std::map output;
+      extract_outputs(output, run_model_canonical(inputs, spin));
+      auto atom_energy = output["atom_energy"]
+                             .reshape({nall_nodes})
+                             .slice(0, 0, nloc)
+                             .contiguous();
+      auto force = output["force"].reshape({nall_nodes, 3}).contiguous();
+      auto force_mag =
+          output["force_mag"].reshape({nall_nodes, 3}).contiguous();
+      // The DPA4C operator, the only one behind this path, returns the force,
+      // the magnetic force, the virial and the atomic energy in fp64, which
+      // ptexpt::copy_results writes into the outputs.
+      std::vector targets = {
+          torch::from_blob(d_force, {nall_nodes, 3}, opt_f64),
+          torch::from_blob(d_force_mag, {nall_nodes, 3}, opt_f64)};
+      std::vector results = {force, force_mag};
+      if (d_virial != nullptr) {
+        targets.push_back(torch::from_blob(d_virial, {9}, opt_f64));
+        results.push_back(output["virial"].reshape({9}));
+      }
+      if (nloc > 0) {
+        targets.push_back(torch::from_blob(d_atom_energy, {nloc}, opt_f64));
+        results.push_back(atom_energy);
+      }
+      ptexpt::copy_results(targets, results, nall_nodes);
+      if (d_atom_virial != nullptr) {
+        torch::from_blob(d_atom_virial, {nall_nodes, 9}, opt_f64)
+            .copy_(output["atom_virial"].reshape({nall_nodes, 9}));
+      }
+    };
+    // A graph without halo nodes keeps its extents from step to step, so its
+    // evaluation is replayed; a domain's halo changes its extent at every
+    // neighbor rebuild, so a graph with halo nodes is evaluated eagerly.
+    if (nloc != nall_nodes) {
+      evaluate();
+    } else {
+      if (!canonical_replay_) {
+        canonical_replay_ = std::make_unique();
+      }
+      canonical_replay_->run(
+          {ptexpt::address_key(d_atom_energy), ptexpt::address_key(d_force),
+           ptexpt::address_key(d_force_mag), ptexpt::address_key(d_virial),
+           ptexpt::address_key(d_atom_virial), ptexpt::address_key(d_atype),
+           ptexpt::address_key(d_source), ptexpt::address_key(d_edge_vec),
+           ptexpt::address_key(d_destination_row_ptr),
+           ptexpt::address_key(source_row_data),
+           ptexpt::address_key(source_order_data), ptexpt::address_key(d_spin),
+           nloc, nall_nodes, edge_storage, index_width},
+          gpu_id, evaluate);
+    }
+    if (stream == nullptr) {
+      synchronize_current_accelerator_stream(gpu_id);
     }
-    torch::from_blob(d_force, {nall_nodes, 3}, opt_f64).copy_(force);
-    torch::from_blob(d_force_mag, {nall_nodes, 3}, opt_f64).copy_(force_mag);
-    torch::from_blob(d_atom_virial, {nall_nodes, 9}, opt_f64)
-        .copy_(atom_virial);
-    synchronize_current_accelerator_stream();
   });
 }
 
diff --git a/source/api_cc/src/commonPTExpt.h b/source/api_cc/src/commonPTExpt.h
index 8b8696f55f..1981531ccf 100644
--- a/source/api_cc/src/commonPTExpt.h
+++ b/source/api_cc/src/commonPTExpt.h
@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 // Shared utilities for pt_expt (.pt2 / AOTInductor) backend classes.
-// Provides: JSON parser, ZIP archive reader, type-sorted nlist builder,
-// and helpers for the with-comm dual-artifact layout.
+// Provides: the caller-stream scope and the CUDA graph replay of
+// device-resident evaluations, JSON parser, ZIP archive reader, type-sorted
+// nlist builder, and helpers for the with-comm dual-artifact layout.
 #pragma once
 
 #include 
@@ -15,21 +16,318 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
 #include 
 #include 
 
+#include "canonical_graph.h"
 #include "common.h"  // for remap_comm_sendlist
 #include "errors.h"
 #include "neighbor_list.h"
 
+#ifdef DEEPMD_TORCH_HAS_CUDA
+#include 
+#include 
+#include 
+#include 
+#include 
+
+#include 
+#endif
+
 namespace deepmd {
 namespace ptexpt {
 
+/** Validate source width before wrapping caller buffers or replaying a graph.
+ */
+inline void check_canonical_index(const void* source,
+                                  DP_CanonicalIndexWidth width,
+                                  std::int64_t node_count) {
+  if (width != DP_CANONICAL_INDEX32 && width != DP_CANONICAL_INDEX64) {
+    throw deepmd_exception("canonical index width must be four or eight bytes");
+  }
+  if (width == DP_CANONICAL_INDEX32 &&
+      canonical_index_width(node_count) != width) {
+    throw deepmd_exception("canonical node count requires 64-bit indices");
+  }
+  if (reinterpret_cast(source) % width != 0) {
+    throw deepmd_exception(
+        "canonical source must be aligned to its word width");
+  }
+}
+
+/**
+ * @brief Make a caller's accelerator stream current for one scope.
+ *
+ * Device-resident entry points take the stream on which the caller produced
+ * their inputs and will consume their outputs. While the scope lives, every
+ * kernel the model launches is enqueued on that stream, so the evaluation is
+ * ordered after the inputs and before the outputs' consumers without a host
+ * synchronization. A null stream leaves the current stream in place.
+ */
+class CallerStreamScope {
+ public:
+  CallerStreamScope(void* stream, int device) {
+    if (stream == nullptr) {
+      return;
+    }
+#ifdef DEEPMD_TORCH_HAS_CUDA
+    guard_.emplace(c10::cuda::getStreamFromExternal(
+        static_cast(stream),
+        static_cast(device)));
+#else
+    (void)device;
+    throw deepmd::deepmd_exception(
+        "a caller stream requires deepmd-kit built with USE_CUDA_TOOLKIT "
+        "against a CUDA-enabled PyTorch.");
+#endif
+  }
+
+ private:
+#ifdef DEEPMD_TORCH_HAS_CUDA
+  std::optional guard_;
+#endif
+};
+
+/**
+ * @brief Lock of the operations of the backend that capture a CUDA graph or
+ * wait for a whole device.
+ *
+ * A wait for a whole device is invalid while any stream of the device
+ * captures, and it invalidates the capture. Every such operation of the
+ * backend holds this lock, so models in different threads of one process
+ * capture, release their graphs and replace their constants in turn. The
+ * lock outlives every model, including one destroyed with the static objects
+ * of the process.
+ */
+inline std::mutex& device_wide_mutex() {
+  static auto* const mutex = new std::mutex;
+  return *mutex;
+}
+
+/**
+ * @brief Wait for the work in flight on every stream of ``device``.
+ *
+ * An evaluation on a caller's stream may still read the constants of a model
+ * while the caller replaces them, so whoever replaces them waits for it first.
+ * A build without caller streams leaves no such work.
+ */
+inline void synchronize_device(const int device) {
+#ifdef DEEPMD_TORCH_HAS_CUDA
+  const std::lock_guard device_wide(device_wide_mutex());
+  const c10::cuda::CUDAGuard device_guard(
+      static_cast(device));
+  c10::cuda::device_synchronize();
+#else
+  (void)device;
+#endif
+}
+
+/// Address of a buffer as an entry of a GraphReplay key.
+inline std::int64_t address_key(const void* pointer) {
+  return static_cast(reinterpret_cast(pointer));
+}
+
+/**
+ * @brief Copy the results of an evaluation of ``nodes`` nodes into the
+ * caller's outputs, ``results[k]`` into ``targets[k]``; both lists hold the
+ * same number of tensors, at least one.
+ *
+ * Up to 8,192 nodes the results take one multi-tensor copy, which copies a
+ * list of results of one dtype in a single kernel. That copy serves every
+ * chunk of 65,536 elements with one block, so on the force of a larger graph
+ * it leaves most of the device idle; a larger graph therefore copies each
+ * tensor on its own, and each copy spreads over the whole device. Inside a
+ * graph replay on compute capability 12.0, the outputs of a DPA4C energy
+ * model take 10 us in one kernel and 6 us in separate copies at 16,384 nodes,
+ * and 18 and 8 us at 100,800.
+ */
+inline void copy_results(const std::vector& targets,
+                         const std::vector& results,
+                         const std::int64_t nodes) {
+  constexpr std::int64_t kMultiTensorNodes = 8192;
+  if (nodes <= kMultiTensorNodes) {
+    at::_foreach_copy_(targets, results);
+    return;
+  }
+  for (std::size_t k = 0; k < targets.size(); ++k) {
+    targets[k].copy_(results[k]);
+  }
+}
+
+/**
+ * @brief Replay repeated device-resident evaluations as one CUDA graph.
+ *
+ * Launching kernels dominates an evaluation of a small system. Two
+ * evaluations over the same buffers, extents and host values launch the same
+ * kernels with the same arguments, so the evaluation is captured into a CUDA
+ * graph on the second consecutive call with the same key -- every buffer
+ * address, extent and host value that shapes its launches -- and relaunched
+ * as a whole while the key stays the same. A call with another key releases
+ * the graph and evaluates eagerly.
+ *
+ * The evaluation must be free of host synchronization, and every device
+ * tensor it uses beyond the caller's buffers must either be allocated inside
+ * it, so that the graph owns them, or persist at an address that the key
+ * covers; a replay repeats the launches of the capture without the host code
+ * that made them, so every host value that shapes them belongs to the key.
+ * The constants of the model are read where they reside at capture, so
+ * whoever replaces them calls ``invalidate`` first. An evaluation that cannot
+ * be captured runs eagerly instead (see ``capture``).
+ *
+ * A capture, like the release of a graph, waits for the whole device, and
+ * both hold ``device_wide_mutex``, so the replays of different models capture
+ * in turn and the evaluations and replays of other threads proceed
+ * meanwhile. CUDA work outside the backend that waits for the whole device
+ * while a capture is under way invalidates that capture, and the replay then
+ * evaluates eagerly from there on.
+ */
+class GraphReplay {
+ public:
+  using Key = std::vector;
+
+  /**
+   * @brief Evaluate under ``key`` on the current stream of ``device``.
+   *
+   * @param key Every buffer address, extent and host value that shapes the
+   * launches of ``evaluate``.
+   * @param device Index of the CUDA device.
+   * @param evaluate The evaluation, enqueued on the current stream.
+   */
+  template 
+  void run(const Key& key, const int device, Evaluate&& evaluate) {
+#ifdef DEEPMD_TORCH_HAS_CUDA
+    if (graph_ && key == key_) {
+      graph_->replay();
+      return;
+    }
+    if (key != key_ || !capturable_) {
+      // The key is recorded once the evaluation has completed, so that a
+      // capture always follows an eager evaluation under the same key, which
+      // has run every one-time step of the program (such as the folding of
+      // its constants) outside the graph.
+      release();
+      key_.clear();
+      evaluate();
+      key_ = key;
+      return;
+    }
+    if (capture(device, evaluate)) {
+      graph_->replay();
+    } else {
+      evaluate();
+    }
+#else
+    (void)key;
+    (void)device;
+    evaluate();
+#endif
+  }
+
+  /**
+   * @brief Release the graph, so that the next call evaluates eagerly.
+   *
+   * Called before the constants that a captured evaluation reads change.
+   */
+  void invalidate() {
+#ifdef DEEPMD_TORCH_HAS_CUDA
+    release();
+    key_.clear();
+#endif
+  }
+
+ private:
+#ifdef DEEPMD_TORCH_HAS_CUDA
+  void release() {
+    if (graph_) {
+      // Replays still in flight run in the private pool of the graph, on the
+      // device it was captured on.
+      const std::lock_guard device_wide(device_wide_mutex());
+      const c10::cuda::CUDAGuard device_guard(
+          static_cast(device_));
+      c10::cuda::device_synchronize();
+      graph_.reset();
+    }
+  }
+
+  /**
+   * @brief Capture ``evaluate`` into a graph, which becomes the graph of the
+   * current key only once its capture has ended.
+   *
+   * A capture that fails at any stage -- its beginning, the evaluation, its
+   * end or the instantiation of the graph -- leaves no graph behind and turns
+   * capturing off for this replay, and the caller evaluates eagerly: an
+   * evaluation that has run eagerly under the same key fails in a capture for
+   * a reason of the capture, which a retry would meet again while holding
+   * another private memory pool. An error of the evaluation itself recurs in
+   * the eager evaluation, which reports it.
+   *
+   * @return Whether the capture succeeded.
+   */
+  template 
+  bool capture(const int device, Evaluate& evaluate) {
+    const std::lock_guard device_wide(device_wide_mutex());
+    // The cached blocks of the eager evaluations are released before the
+    // graph allocates its private pool, so the pool takes their place in
+    // device memory instead of adding to it.
+    c10::cuda::device_synchronize();
+    c10::cuda::CUDACachingAllocator::emptyCache();
+    auto graph = std::make_unique();
+    try {
+      // A capture needs a stream other than the legacy default one, on which
+      // a caller may evaluate; the graph replays on the caller's stream.
+      const c10::cuda::CUDAStreamGuard capture_stream(
+          c10::cuda::getStreamFromPool(false,
+                                       static_cast(device)));
+      graph->capture_begin({0, 0}, cudaStreamCaptureModeThreadLocal);
+      try {
+        evaluate();
+      } catch (...) {
+        // End the capture the failure interrupted, so that the stream leaves
+        // capture mode; the failure of the evaluation is the one reported.
+        try {
+          graph->capture_end();
+        } catch (const c10::Error&) {
+        }
+        throw;
+      }
+      graph->capture_end();
+    } catch (const std::exception& error) {
+      // The graph of a failed capture is kept for the lifetime of the
+      // replay, since the caching allocator may still refer to it, and the
+      // memory its partial evaluation allocated may stay reserved in its
+      // private pool for the rest of the process. Capturing is off from here
+      // on, so at most one is kept.
+      abandoned_ = std::move(graph);
+      capturable_ = false;
+      const auto* const c10_error = dynamic_cast(&error);
+      std::cerr << "WARNING: deepmd-kit could not capture the evaluation into "
+                   "a CUDA graph ("
+                << (c10_error != nullptr ? c10_error->what_without_backtrace()
+                                         : error.what())
+                << "); later evaluations run without graph replay."
+                << std::endl;
+      return false;
+    }
+    graph_ = std::move(graph);
+    device_ = device;
+    return true;
+  }
+
+  Key key_;
+  int device_ = 0;
+  bool capturable_ = true;
+  std::unique_ptr graph_;
+  std::unique_ptr abandoned_;
+#endif
+};
+
 // ============================================================================
 // Minimal JSON value parser for reading metadata from .pt2 archives.
 // Supports: strings, numbers, booleans, arrays, objects.
@@ -941,22 +1239,20 @@ class ChargeStateFold {
           " tables but the archive names " + std::to_string(constants_.size()) +
           " constants; it cannot serve a runtime charge state");
     }
-    // The inactive buffer is a complete model image. A partial inactive update
-    // copies tensor constants and buffers but omits ordinary parameters, so it
-    // cannot be swapped into service safely.
-    auto* runner = target.get_runner();
-    auto constants = runner->extract_constants_map(/*use_inactive=*/false);
+    // The tables are copied into the active constants in place, which keeps
+    // their addresses and marks the constants for folding again at the next
+    // run. The caller waits for the device on both sides of the update, so no
+    // evaluation reads the constants while they change.
+    std::unordered_map updated;
     for (size_t ii = 0; ii < tables.size(); ++ii) {
       // An unnamed output belongs to a mechanism this model has disabled and
       // has no constant to reach.
       if (!constants_[ii].empty()) {
-        constants[constants_[ii]] = tables[ii];
+        updated.emplace(constants_[ii], tables[ii]);
       }
     }
-    target.load_constants(constants, /*use_inactive=*/true,
-                          /*check_full_update=*/true);
-    runner->swap_constant_buffer();
-    runner->free_inactive_constant_buffer();
+    target.load_constants(updated, /*use_inactive=*/false,
+                          /*check_full_update=*/false);
   }
 
  private:
diff --git a/source/api_cc/src/longRangePTExpt.h b/source/api_cc/src/longRangePTExpt.h
new file mode 100644
index 0000000000..65c3c21b6e
--- /dev/null
+++ b/source/api_cc/src/longRangePTExpt.h
@@ -0,0 +1,298 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+#pragma once
+
+#include 
+
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+
+#include "commonPT.h"
+
+#ifdef DEEPMD_TORCH_HAS_CUDA
+#include 
+#include 
+#include 
+#include 
+#endif
+
+namespace deepmd {
+namespace ptexpt {
+
+/**
+ * @brief Build the all-pairs topology a long-range module of the graph lower
+ * reads.
+ *
+ * Inference-path twin of Python ``build_pair_graph`` in
+ * ``deepmd/dpmodel/utils/neighbor_graph/pair_graph.py`` for a single frame:
+ * every unordered pair of nodes once, ``[src, dst]`` with ``src < dst``, its
+ * vector ``r_src - r_dst``, and two masked guard pairs that keep the pair axis
+ * at the lower bound of the exported lower. The model-level pair exclusion is
+ * applied to the mask, as it is to the neighbor graph.
+ *
+ * @param coord Node positions, shape ``(N, 3)``, float64, on the model device.
+ * @param atype Node types, shape ``(N,)``, int64.
+ * @param type_mask_table Device keep table of the pair exclusion, or an
+ * undefined tensor when the model excludes no pair.
+ * @param ntypes Number of atom types.
+ * @param vec_dtype Precision of the pair vectors, that of the edge vectors of
+ * the artifact.
+ * @return ``{pair_index (2, P + 2) int64, pair_vec (P + 2, 3), pair_mask
+ * (P + 2,) bool}``.
+ */
+inline std::vector buildAllPairsTopology(
+    const torch::Tensor& coord,
+    const torch::Tensor& atype,
+    const torch::Tensor& type_mask_table,
+    const int ntypes,
+    const torch::ScalarType vec_dtype) {
+  const std::int64_t node_count = coord.size(0);
+  const auto index_options =
+      torch::TensorOptions().dtype(torch::kInt64).device(coord.device());
+  const auto mask_options =
+      torch::TensorOptions().dtype(torch::kBool).device(coord.device());
+  const torch::Tensor index =
+      torch::triu_indices(node_count, node_count, 1, index_options);
+  const torch::Tensor vec = (coord.index_select(0, index.select(0, 0)) -
+                             coord.index_select(0, index.select(0, 1)))
+                                .to(vec_dtype);
+  const torch::Tensor mask =
+      applyPairExclusion(index, torch::ones({index.size(1)}, mask_options),
+                         atype, type_mask_table, ntypes);
+  return {torch::cat({index, torch::zeros({2, 2}, index_options)}, 1),
+          torch::cat({vec, torch::zeros({2, 3}, vec.options())}, 0),
+          torch::cat({mask, torch::zeros({2}, mask_options)})};
+}
+
+/**
+ * @brief Build the all-pairs topology of a periodic frame, which holds no
+ * pair: the reciprocal-space topology describes the interactions.
+ *
+ * Counterpart of Python ``empty_pair_graph`` for the exported lower: two
+ * masked guard pairs, the lower bound of its pair axis. Python pads four,
+ * which keeps the pair axis of a traced sample apart from the axes of length
+ * two and three; the exported lower accepts any length from two on.
+ */
+inline std::vector buildGuardPairTopology(
+    const torch::Device& device, const torch::ScalarType vec_dtype) {
+  const auto options = torch::TensorOptions().device(device);
+  return {torch::zeros({2, 2}, options.dtype(torch::kInt64)),
+          torch::zeros({2, 3}, options.dtype(vec_dtype)),
+          torch::zeros({2}, options.dtype(torch::kBool))};
+}
+
+/**
+ * @brief The smallest even integer of at least ``n`` without a prime factor
+ * above 7.
+ */
+inline std::int64_t evenFftSize(std::int64_t n) {
+  std::int64_t m = std::max(2, n + n % 2);
+  while (true) {
+    std::int64_t rest = m;
+    for (const std::int64_t p : {2, 3, 5, 7}) {
+      while (rest % p == 0) {
+        rest /= p;
+      }
+    }
+    if (rest == 1) {
+      return m;
+    }
+    m += 2;
+  }
+}
+
+/**
+ * @brief The size of the mesh of a periodic frame along its lattice vectors.
+ *
+ * Twin of Python ``mesh_dims_of_cells`` and of the long-range operator for a
+ * single frame: along every lattice vector ``a_d`` the mesh holds the
+ * smallest even number of points without a prime factor above 7 of at least
+ * ``ceil(|a_d| / spacing)``. The squared length sums the squares of the
+ * components, each rounded on its own as Python forms it; the explicit fma of
+ * a square with zero is that rounded square, which the compiler does not
+ * fuse into the sum, so that the three size the mesh identically whatever
+ * the floating-point contraction of the build.
+ *
+ * @param box The row-major cell, the rows its lattice vectors, nine values in
+ * the length unit of the model.
+ * @param spacing The largest mesh spacing, in the length unit of the model.
+ * @throw deepmd::deepmd_exception If the mesh exceeds 2^20 points along a
+ * lattice vector, the limit of the operator.
+ */
+inline std::array meshDims(const std::vector& box,
+                                            const double spacing) {
+  std::array dims = {};
+  for (int d = 0; d < 3; ++d) {
+    const double* v = box.data() + 3 * d;
+    const double length =
+        std::sqrt(std::fma(v[0], v[0], 0.0) + std::fma(v[1], v[1], 0.0) +
+                  std::fma(v[2], v[2], 0.0));
+    const double points = std::ceil(length / spacing);
+    if (!(points < double(1 << 20))) {
+      throw deepmd::deepmd_exception(
+          "the mesh of the cell exceeds 2^20 points along a lattice vector; "
+          "a lattice vector of " +
+          std::to_string(length) + " at a mesh spacing of " +
+          std::to_string(spacing) + " needs " + std::to_string(points) + ".");
+    }
+    dims[d] = evenFftSize(static_cast(points));
+  }
+  return dims;
+}
+
+/**
+ * @brief Build the mesh topology a periodic long-range module of the graph
+ * lower reads.
+ *
+ * Twin of Python ``build_mesh_grid`` in
+ * ``deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py`` for a single
+ * frame, on the mesh of ``meshDims``. Without a cell, the frame is
+ * non-periodic: the cell is the identity, the mesh minimal and the frame
+ * masked.
+ *
+ * @param coord Node positions, shape ``(N, 3)``, on the model device.
+ * @param box The row-major cell, the rows its lattice vectors, nine values in
+ * the length unit of the model, or empty for a non-periodic frame.
+ * @param spacing The largest mesh spacing, in the length unit of the model.
+ * @param dtype Precision of the geometry, that of the edge vectors.
+ * @return ``{box (1, 3, 3), volume (1,), coord (N, 3), strain (1, 3, 3),
+ * mask (1,) bool, mode_1 (n1,), mode_2 (n2,), mode_3 (n3,)}``.
+ */
+inline std::vector buildMeshTopology(
+    const torch::Tensor& coord,
+    const std::vector& box,
+    const double spacing,
+    const torch::ScalarType dtype) {
+  const auto options =
+      torch::TensorOptions().dtype(dtype).device(coord.device());
+  const bool periodic = !box.empty();
+  std::array cell = {1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0};
+  std::array dims = {2, 2, 2};
+  if (periodic) {
+    for (int k = 0; k < 9; ++k) {
+      cell[k] = box[k];
+    }
+    dims = meshDims(box, spacing);
+  }
+  const double volume =
+      std::abs(cell[0] * (cell[4] * cell[8] - cell[5] * cell[7]) -
+               cell[1] * (cell[3] * cell[8] - cell[5] * cell[6]) +
+               cell[2] * (cell[3] * cell[7] - cell[4] * cell[6]));
+  // The copy detaches the tensor from the local array, which a conversion
+  // to float64 on the host would alias.
+  const auto cell_tensor =
+      torch::from_blob(cell.data(), {1, 3, 3}, torch::kFloat64)
+          .to(options, /*non_blocking=*/false, /*copy=*/true);
+  std::vector fields = {
+      cell_tensor, torch::full({1}, volume, options), coord.to(dtype),
+      torch::zeros({1, 3, 3}, options),
+      torch::full({1}, periodic, options.dtype(torch::kBool))};
+  for (const std::int64_t n : dims) {
+    // The integer frequencies of an axis of n points, in FFT order.
+    const auto m = torch::arange(n, options);
+    fields.push_back(torch::where(m > n / 2, m - n, m));
+  }
+  return fields;
+}
+
+/**
+ * @brief The cell input of the device-resident evaluations of a model with a
+ * long-range module, in pinned host memory.
+ *
+ * The compact lower reads the cell on the device while the stream runs its
+ * evaluation, so that a CUDA graph capturing the evaluation reads the cell
+ * of every replay from the same memory. One event orders those reads: an
+ * evaluation on another stream than the previous one first waits on the
+ * device for the previous one, so that the event of the last evaluation
+ * follows every earlier read, and a cell other than the one held waits for
+ * that event before the memory is written.
+ */
+class PinnedCell {
+ public:
+  PinnedCell() = default;
+  PinnedCell(const PinnedCell&) = delete;
+  PinnedCell& operator=(const PinnedCell&) = delete;
+  // The destructor reports no error.
+  ~PinnedCell() {
+#ifdef DEEPMD_TORCH_HAS_CUDA
+    if (read_ != nullptr) {
+      int current = 0;
+      cudaGetDevice(¤t);
+      cudaSetDevice(device_);
+      cudaEventDestroy(read_);
+      cudaSetDevice(current);
+    }
+#endif
+  }
+
+  /**
+   * @brief Hold ``box`` for the next evaluation, which is enqueued on the
+   * current stream of ``device``.
+   *
+   * @param box The row-major cell, nine values, or empty for a non-periodic
+   * frame, whose cell is zero.
+   * @param device Index of the CUDA device of the evaluation.
+   * @return The cell input, a pinned float64 CPU tensor of shape ``(1, 9)``.
+   */
+  const torch::Tensor& hold(const std::vector& box, const int device) {
+    if (!cell_.defined()) {
+      cell_ = torch::zeros(
+          {1, 9},
+          torch::TensorOptions().dtype(torch::kFloat64).pinned_memory(true));
+    }
+#ifdef DEEPMD_TORCH_HAS_CUDA
+    device_ = static_cast(device);
+    stream_ = c10::cuda::getCurrentCUDAStream(device_).stream();
+    if (read_ != nullptr && stream_ != read_stream_) {
+      C10_CUDA_CHECK(cudaStreamWaitEvent(stream_, read_, 0));
+    }
+#else
+    (void)device;
+#endif
+    std::array values = {};
+    std::copy(box.begin(), box.end(), values.begin());
+    double* const cell = cell_.data_ptr();
+    if (std::memcmp(cell, values.data(), sizeof(values)) != 0) {
+#ifdef DEEPMD_TORCH_HAS_CUDA
+      if (read_ != nullptr) {
+        C10_CUDA_CHECK(cudaEventSynchronize(read_));
+      }
+#endif
+      std::memcpy(cell, values.data(), sizeof(values));
+    }
+    return cell_;
+  }
+
+  /**
+   * @brief Mark the evaluation of the cell last held, once enqueued, as the
+   * last that reads the cell.
+   */
+  void record() {
+#ifdef DEEPMD_TORCH_HAS_CUDA
+    const c10::cuda::CUDAGuard device_guard(device_);
+    if (read_ == nullptr) {
+      C10_CUDA_CHECK(cudaEventCreateWithFlags(&read_, cudaEventDisableTiming));
+    }
+    C10_CUDA_CHECK(cudaEventRecord(read_, stream_));
+    read_stream_ = stream_;
+#endif
+  }
+
+ private:
+  torch::Tensor cell_;
+#ifdef DEEPMD_TORCH_HAS_CUDA
+  c10::DeviceIndex device_ = 0;
+  // The stream of the evaluation of the cell last held.
+  cudaStream_t stream_ = nullptr;
+  // The last evaluation that reads the cell, and its stream.
+  cudaEvent_t read_ = nullptr;
+  cudaStream_t read_stream_ = nullptr;
+#endif
+};
+
+}  // namespace ptexpt
+}  // namespace deepmd
diff --git a/source/api_cc/tests/CMakeLists.txt b/source/api_cc/tests/CMakeLists.txt
index 1d1b817957..1dbbc962f6 100644
--- a/source/api_cc/tests/CMakeLists.txt
+++ b/source/api_cc/tests/CMakeLists.txt
@@ -20,6 +20,11 @@ if(ENABLE_PYTORCH)
   # BUILD_PT_EXPT is set for the test binary; otherwise pt_expt tests all
   # GTEST_SKIP() with "PyTorch support is not enabled".
   target_link_libraries(runUnitTests_cc "${TORCH_LIBRARIES}")
+  # The device-resident cases order their work on a caller's stream through the
+  # c10 CUDA stream API, which the backend uses under the same condition.
+  if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA)
+    target_compile_definitions(runUnitTests_cc PRIVATE DEEPMD_TORCH_HAS_CUDA)
+  endif()
 endif()
 if(ENABLE_JAX)
   target_compile_definitions(runUnitTests_cc PRIVATE BUILD_JAX)
diff --git a/source/api_cc/tests/dpa4c_canonical_test_common.h b/source/api_cc/tests/dpa4c_canonical_test_common.h
new file mode 100644
index 0000000000..4d6fa22d81
--- /dev/null
+++ b/source/api_cc/tests/dpa4c_canonical_test_common.h
@@ -0,0 +1,829 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Cases of the device-resident entry point of the compressed DPA4C archives
+// frozen with the compact canonical lower
+// (source/tests/infer/gen_dpa4c_canonical.py): an energy model, evaluated
+// through a DeepPot interface, and a native-spin model, evaluated through a
+// DeepSpin interface. The C++ suite instantiates them for deepmd::DeepPot and
+// deepmd::DeepSpin, the C suite for deepmd::hpp::DeepPot and
+// deepmd::hpp::DeepSpin and so for DP_DPA4ComputeCanonicalGraphGPU and
+// DP_DPA4SpinComputeCanonicalGraphGPU. The including file provides
+// _build_nlist through the test_utils.h of its suite.
+//
+// compute_canonical_graph_gpu evaluates a graph whose buffers, like its
+// outputs, reside on the model device. The cases hand it the graph that the
+// host entry points build from the same coordinates, so the two entry points
+// evaluate one program on equal inputs and agree exactly. The graph of one
+// rank has no halo nodes, so the repeated evaluations of a case run eagerly,
+// are captured into a CUDA graph and are replayed. Every evaluation writes
+// into outputs reset to NaN, and no case may fall back from the replay to
+// eager evaluation, so the cases check the results of the replay itself: on
+// the current stream and on a caller's stream, without the optional outputs
+// and inputs, after its inputs change in place, across a change of its key,
+// while another thread holds a capture of its own open, and under a charge
+// state applied at run time. A graph with halo nodes, which domain
+// decomposition hands over, is evaluated eagerly instead.
+//
+// The cases form a typed test suite in this namespace, which each suite
+// instantiates in the same namespace.
+#pragma once
+
+#include 
+
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+
+#include "DeepPotPTExpt.h"
+#include "expected_ref.h"
+
+#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT
+#include 
+
+#include "commonPT.h"
+#ifdef DEEPMD_TORCH_HAS_CUDA
+#include 
+#include 
+#include 
+#include 
+
+#include 
+#include 
+#include 
+#endif
+#endif
+
+namespace deepmd_test {
+namespace dpa4c_canonical {
+
+// The charge state applied at run time; the archives are frozen with another.
+const std::vector kOtherChargeSpin = {1.0, 2.0};
+
+// The fixtures are fp32 inside; the references come from the Python
+// inference of the same archives on graphs of another edge order.
+constexpr double kReferenceTolerance = 1e-5;
+
+// Uniform strain of the system between two evaluations of one graph. Every
+// atom keeps its fractional coordinates, and so its cell of the neighbor
+// search, and every distance changes by far less than its margin to the
+// cutoff, so the graph keeps its edges and their order while every edge
+// vector changes.
+constexpr double kStrain = 1e-4;
+
+// Cycles a caller's stream is held back before an evaluation is enqueued on
+// it: tens of milliseconds, far longer than the evaluation.
+constexpr std::int64_t kHoldCycles = std::int64_t{1} << 27;
+
+// The warning of a replay whose capture failed, which evaluates eagerly from
+// then on.
+constexpr const char* kCaptureWarning = "could not capture the evaluation";
+
+/// The energy model, evaluated through the DeepPot interface ``DP``.
+template 
+struct Potential {
+  using Evaluator = DP;
+  static constexpr bool kSpin = false;
+  static constexpr const char* kName = "Potential";
+  static constexpr const char* kModelPath =
+      "../../tests/infer/deeppot_dpa4c_canonical.pt2";
+  static constexpr const char* kRefPath =
+      "../../tests/infer/deeppot_dpa4c_canonical.expected";
+};
+
+/// The native-spin model, evaluated through the DeepSpin interface ``DS``.
+template 
+struct Spin {
+  using Evaluator = DS;
+  static constexpr bool kSpin = true;
+  static constexpr const char* kName = "Spin";
+  static constexpr const char* kModelPath =
+      "../../tests/infer/deeppot_dpa4c_canonical_spin.pt2";
+  static constexpr const char* kRefPath =
+      "../../tests/infer/deeppot_dpa4c_canonical_spin.expected";
+};
+
+/// Names the instantiations of the cases by their model.
+struct CaseNames {
+  template 
+  static std::string GetName(int) {
+    return Case::kName;
+  }
+};
+
+/// A configuration of the system: coordinates, magnetic moments (empty for
+/// the energy model), atom types and cell.
+struct Configuration {
+  std::vector coord;
+  std::vector spin;
+  std::vector atype;
+  std::vector box;
+};
+
+/// Results of one evaluation, on the host. An output the evaluation was not
+/// asked for, or that the model does not produce, is empty.
+struct Results {
+  std::vector atom_energy;
+  std::vector force;
+  std::vector force_mag;
+  std::vector virial;
+  std::vector atom_virial;
+};
+
+/// Largest absolute difference of two equally long arrays, NaN when an entry
+/// of either is NaN.
+inline double max_difference(const std::vector& actual,
+                             const std::vector& expected) {
+  EXPECT_EQ(actual.size(), expected.size());
+  double difference = 0.0;
+  for (std::size_t ii = 0; ii < std::min(actual.size(), expected.size());
+       ++ii) {
+    const double entry = std::abs(actual[ii] - expected[ii]);
+    if (std::isnan(entry)) {
+      return entry;
+    }
+    difference = std::max(difference, entry);
+  }
+  return difference;
+}
+
+/// Largest magnitude of an array.
+inline double largest(const std::vector& values) {
+  double magnitude = 0.0;
+  for (const double value : values) {
+    magnitude = std::max(magnitude, std::abs(value));
+  }
+  return magnitude;
+}
+
+/// Sum the per-node rows of ``values``, three per node, onto the atoms that
+/// own the nodes.
+inline std::vector fold_onto_owners(const std::vector& values,
+                                            const std::vector& owner,
+                                            const int natoms) {
+  std::vector folded(3 * static_cast(natoms), 0.0);
+  for (std::size_t node = 0; node < owner.size(); ++node) {
+    for (int dd = 0; dd < 3; ++dd) {
+      folded[3 * owner[node] + dd] += values[3 * node + dd];
+    }
+  }
+  return folded;
+}
+
+#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT && defined(DEEPMD_TORCH_HAS_CUDA)
+/// A CUDA graph capture that another thread holds open on a stream of its own
+/// from construction until ``close``.
+class CaptureElsewhere {
+ public:
+  CaptureElsewhere() : thread_([this] { hold(); }) {
+    std::unique_lock lock(mutex_);
+    changed_.wait(lock, [this] { return open_; });
+  }
+  ~CaptureElsewhere() { close(); }
+
+  /// End the capture and return the status with which it ended:
+  /// ``cudaSuccess`` unless something invalidated it while it was open.
+  cudaError_t close() {
+    {
+      const std::lock_guard lock(mutex_);
+      closing_ = true;
+    }
+    changed_.notify_all();
+    if (thread_.joinable()) {
+      thread_.join();
+    }
+    return status_;
+  }
+
+ private:
+  void hold() {
+    const c10::cuda::CUDAStream stream =
+        c10::cuda::getStreamFromPool(/*isHighPriority=*/false, /*device=*/0);
+    const cudaError_t begun = cudaStreamBeginCapture(
+        stream.stream(), cudaStreamCaptureModeThreadLocal);
+    {
+      const std::lock_guard lock(mutex_);
+      open_ = true;
+    }
+    changed_.notify_all();
+    {
+      std::unique_lock lock(mutex_);
+      changed_.wait(lock, [this] { return closing_; });
+    }
+    cudaGraph_t graph = nullptr;
+    status_ = begun == cudaSuccess
+                  ? cudaStreamEndCapture(stream.stream(), &graph)
+                  : begun;
+    if (graph != nullptr) {
+      cudaGraphDestroy(graph);
+    }
+    // A failed capture leaves its error for the next query of this thread.
+    cudaGetLastError();
+  }
+
+  std::mutex mutex_;
+  std::condition_variable changed_;
+  bool open_ = false;
+  bool closing_ = false;
+  cudaError_t status_ = cudaSuccess;
+  // Declared last, so that the thread starts after the members it uses.
+  std::thread thread_;
+};
+#endif
+
+/// The device-resident entry point of the model of ``Case``.
+template 
+class TestInferDpa4cCanonicalPtExpt : public ::testing::Test {
+ protected:
+  using Evaluator = typename Case::Evaluator;
+
+  Configuration system;
+  int natoms = 0;
+  Evaluator dp;
+  ExpectedRef ref;
+  Results host;
+  bool capturing_stderr = false;
+
+#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT
+  /// A compact canonical graph on the device, with the owner of each node
+  /// and, for the spin model, the moment of each node.
+  struct Graph {
+    deepmd::CanonicalGraphTensorPack tensors;
+    torch::Tensor spin;
+    std::vector owner;
+  };
+
+  // The compact canonical graph of the host entry points, on the device.
+  torch::Tensor graph_atype;
+  torch::Tensor graph_source;
+  torch::Tensor graph_edge_vec;
+  torch::Tensor graph_destination_row_ptr;
+  torch::Tensor graph_source_row_ptr;
+  torch::Tensor graph_source_order;
+  torch::Tensor graph_spin;
+  int nodes = 0;
+  std::int64_t edge_storage = 0;
+  // Outputs reused by the evaluations of a case, whose addresses the replay
+  // of repeated evaluations keys on.
+  torch::Tensor out_atom_energy;
+  torch::Tensor out_force;
+  torch::Tensor out_force_mag;
+  torch::Tensor out_virial;
+  torch::Tensor out_atom_virial;
+#endif
+
+  void SetUp() override {
+#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT
+    GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#else
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "Skip because the compact canonical path needs a CUDA "
+                      "device.";
+    }
+    std::ifstream model_file(Case::kModelPath);
+    if (!model_file.good()) {
+      GTEST_SKIP() << "Skip because " << Case::kModelPath
+                   << " was not generated (run "
+                      "source/tests/infer/gen_dpa4c_canonical.py).";
+    }
+    dp.init(Case::kModelPath);
+    ASSERT_TRUE(dp.uses_canonical_graph_inference());
+    if constexpr (!Case::kSpin) {
+      ASSERT_FALSE(dp.has_long_range());
+    }
+    ref.load(Case::kRefPath);
+    system.coord = ref.get("system", "coord");
+    system.box = ref.get("system", "box");
+    if constexpr (Case::kSpin) {
+      system.spin = ref.get("system", "spin");
+    }
+    for (const double type : ref.get("system", "atype")) {
+      system.atype.push_back(static_cast(type));
+    }
+    natoms = static_cast(system.atype.size());
+    host = compute_host(dp, system);
+
+    use_graph(build_graph(system, /*fold=*/true));
+    // One rank folds the ghost atoms onto their owners: no halo node.
+    ASSERT_EQ(nodes, natoms);
+    allocate_outputs();
+    // Some cases evaluate on a caller's stream, which is ordered after none
+    // of the work that built the graph.
+    torch::cuda::synchronize();
+    // A replay whose capture fails evaluates eagerly instead, with the same
+    // results; only its warning tells the two apart.
+    testing::internal::CaptureStderr();
+    capturing_stderr = true;
+#endif
+  }
+
+  void TearDown() override {
+    if (!capturing_stderr) {
+      return;
+    }
+    const std::string log = testing::internal::GetCapturedStderr();
+    std::cerr << log;
+    EXPECT_EQ(log.find(kCaptureWarning), std::string::npos)
+        << "an evaluation fell back from the replay to eager evaluation";
+  }
+
+  /// Results of the host entry points of ``model`` for a configuration.
+  static Results compute_host(Evaluator& model, const Configuration& config) {
+    Results results;
+    double energy = 0.0;
+    if constexpr (Case::kSpin) {
+      model.compute(energy, results.force, results.force_mag, results.virial,
+                    results.atom_energy, results.atom_virial, config.coord,
+                    config.spin, config.atype, config.box);
+    } else {
+      model.compute(energy, results.force, results.virial, results.atom_energy,
+                    results.atom_virial, config.coord, config.atype,
+                    config.box);
+    }
+    return results;
+  }
+
+  /// Agreement of the host entry points with the references of a section.
+  void expect_reference(const Results& results, const std::string& section) {
+    EXPECT_LT(max_difference(results.atom_energy,
+                             ref.get(section, "expected_e")),
+              kReferenceTolerance)
+        << section;
+    EXPECT_LT(
+        max_difference(results.force, ref.get(section, "expected_f")),
+        kReferenceTolerance)
+        << section;
+    EXPECT_LT(max_difference(results.virial,
+                             ref.get(section, "expected_tot_v")),
+              kReferenceTolerance)
+        << section;
+    EXPECT_LT(max_difference(results.atom_virial,
+                             ref.get(section, "expected_atom_v")),
+              kReferenceTolerance)
+        << section;
+    if constexpr (Case::kSpin) {
+      EXPECT_LT(max_difference(results.force_mag,
+                               ref.get(section, "expected_fm")),
+                kReferenceTolerance)
+          << section;
+    }
+  }
+
+  /// Exact agreement of an evaluation with the host entry points over the
+  /// outputs it was asked for. NaN, the value of an unwritten output, never
+  /// agrees.
+  static void expect_equal(const Results& actual,
+                           const Results& expected,
+                           const std::string& what) {
+    EXPECT_EQ(actual.atom_energy, expected.atom_energy)
+        << what << ": atom energy";
+    EXPECT_EQ(actual.force, expected.force) << what << ": force";
+    EXPECT_EQ(actual.force_mag, expected.force_mag)
+        << what << ": magnetic force";
+    if (!actual.virial.empty()) {
+      EXPECT_EQ(actual.virial, expected.virial) << what << ": virial";
+    }
+    if (!actual.atom_virial.empty()) {
+      EXPECT_EQ(actual.atom_virial, expected.atom_virial)
+          << what << ": atom virial";
+    }
+  }
+
+#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT
+  /// The compact canonical graph that the host entry points build for a
+  /// configuration, on the device. A folded graph holds the local atoms
+  /// alone, as one rank builds it; an unfolded one gives every ghost atom a
+  /// node of its own, as domain decomposition does, which carries the moment
+  /// of its owner.
+  Graph build_graph(const Configuration& config, const bool fold) const {
+    std::vector> nlist;
+    std::vector coord_cpy;
+    std::vector atype_cpy;
+    std::vector mapping;
+    _build_nlist(nlist, coord_cpy, atype_cpy, mapping, config.coord,
+                         config.atype, config.box,
+                         static_cast(dp.cutoff()));
+    const std::vector mapping_64(mapping.begin(), mapping.end());
+    const torch::Device device(torch::kCUDA, 0);
+    deepmd::GraphTensorPack graph = deepmd::buildGraphTensors(
+        nlist, coord_cpy, atype_cpy, mapping_64, natoms,
+        static_cast(coord_cpy.size() / 3), dp.cutoff(), device,
+        /*row_centers=*/nullptr, fold);
+    deepmd::canonicalizeGraphPayload(graph, graph.atype.size(0));
+    graph.edge_vec = graph.edge_vec.to(torch::kFloat32);
+    Graph result;
+    result.tensors = deepmd::compactCanonicalGraph(graph);
+    deepmd::CanonicalGraphTensorPack& compact = result.tensors;
+    compact.atype = compact.atype.to(torch::kInt64).contiguous();
+    compact.source = compact.source.contiguous();
+    compact.edge_vec = compact.edge_vec.contiguous();
+    compact.destination_row_ptr = compact.destination_row_ptr.contiguous();
+    compact.source_row_ptr = compact.source_row_ptr.contiguous();
+    compact.source_order = compact.source_order.contiguous();
+    result.owner = mapping;
+    result.owner.resize(compact.atype.size(0));
+    if constexpr (Case::kSpin) {
+      std::vector moments(3 * result.owner.size());
+      for (std::size_t node = 0; node < result.owner.size(); ++node) {
+        for (int dd = 0; dd < 3; ++dd) {
+          moments[3 * node + dd] =
+              static_cast(config.spin[3 * result.owner[node] + dd]);
+        }
+      }
+      result.spin =
+          torch::from_blob(moments.data(),
+                           {static_cast(result.owner.size()), 3},
+                           torch::TensorOptions().dtype(torch::kFloat32))
+              .to(device);
+    }
+    return result;
+  }
+
+  /// Evaluate ``graph`` from here on.
+  void use_graph(const Graph& graph) {
+    graph_atype = graph.tensors.atype;
+    graph_source = graph.tensors.source;
+    graph_edge_vec = graph.tensors.edge_vec;
+    graph_destination_row_ptr = graph.tensors.destination_row_ptr;
+    graph_source_row_ptr = graph.tensors.source_row_ptr;
+    graph_source_order = graph.tensors.source_order;
+    graph_spin = graph.spin;
+    nodes = static_cast(graph_atype.size(0));
+    edge_storage = graph_source.size(0);
+  }
+
+  /// Allocate outputs for the graph in use, at new addresses while the
+  /// previous ones are held elsewhere.
+  void allocate_outputs() {
+    const auto f64 = torch::TensorOptions()
+                         .dtype(torch::kFloat64)
+                         .device(torch::Device(torch::kCUDA, 0));
+    out_atom_energy = torch::empty({natoms}, f64);
+    out_force = torch::empty({nodes, 3}, f64);
+    if constexpr (Case::kSpin) {
+      out_force_mag = torch::empty({nodes, 3}, f64);
+    }
+    out_virial = torch::empty({9}, f64);
+    out_atom_virial = torch::empty({nodes, 9}, f64);
+  }
+
+  /// Rewrite the inputs of the graph in place with those of another
+  /// configuration, whose graph has the same edges in the same order.
+  void rewrite_in_place(const Configuration& config) {
+    const Graph other = build_graph(config, /*fold=*/true);
+    // The unsigned indices are compared as int64, which every comparison
+    // kernel supports.
+    const auto same = [](const torch::Tensor& left,
+                         const torch::Tensor& right) {
+      return torch::equal(left.to(torch::kInt64), right.to(torch::kInt64));
+    };
+    ASSERT_TRUE(same(other.tensors.source, graph_source));
+    ASSERT_TRUE(
+        same(other.tensors.destination_row_ptr, graph_destination_row_ptr));
+    ASSERT_TRUE(same(other.tensors.source_row_ptr, graph_source_row_ptr));
+    ASSERT_TRUE(same(other.tensors.source_order, graph_source_order));
+    graph_edge_vec.copy_(other.tensors.edge_vec);
+    graph_atype.copy_(other.tensors.atype);
+    if constexpr (Case::kSpin) {
+      graph_spin.copy_(other.spin);
+    }
+  }
+
+  /// Reset every output to NaN on the current stream, so that an evaluation
+  /// that leaves an output unwritten fails the comparison.
+  void reset_outputs() {
+    const double nan = std::numeric_limits::quiet_NaN();
+    out_atom_energy.fill_(nan);
+    out_force.fill_(nan);
+    if constexpr (Case::kSpin) {
+      out_force_mag.fill_(nan);
+    }
+    out_virial.fill_(nan);
+    out_atom_virial.fill_(nan);
+  }
+
+  /// Enqueue an evaluation of the graph on ``stream``, or on the current
+  /// stream when it is null. Without the source CSR the per-atom virial is
+  /// declined.
+  void run_device_entry(void* stream,
+                        const bool source_csr = true,
+                        const bool virial = true) {
+    double* const atom_energy = out_atom_energy.data_ptr();
+    double* const force = out_force.data_ptr();
+    double* const total_virial =
+        virial ? out_virial.data_ptr() : nullptr;
+    double* const atom_virial =
+        source_csr ? out_atom_virial.data_ptr() : nullptr;
+    const auto* const atype = graph_atype.data_ptr();
+    const auto* const source = graph_source.data_ptr();
+    const float* const edge_vec = graph_edge_vec.data_ptr();
+    const auto* const destination_row_ptr =
+        graph_destination_row_ptr.data_ptr();
+    const std::int64_t* const source_row_ptr =
+        source_csr ? graph_source_row_ptr.data_ptr() : nullptr;
+    const std::int64_t* const source_order =
+        source_csr ? graph_source_order.data_ptr() : nullptr;
+    if constexpr (Case::kSpin) {
+      dp.compute_canonical_graph_gpu(
+          atom_energy, force, out_force_mag.data_ptr(), total_virial,
+          atom_virial, atype, source, deepmd::canonical_index_width(nodes),
+          edge_vec, destination_row_ptr, source_row_ptr, source_order,
+          graph_spin.data_ptr(), natoms, nodes, edge_storage, stream);
+    } else {
+      dp.compute_canonical_graph_gpu(
+          atom_energy, force, total_virial, atom_virial, atype, source,
+          deepmd::canonical_index_width(nodes), edge_vec, destination_row_ptr,
+          source_row_ptr, source_order,
+          /*d_coord=*/nullptr, /*world=*/nullptr, /*box=*/nullptr, natoms,
+          nodes, edge_storage, stream);
+    }
+  }
+
+  /// The outputs an evaluation was asked for, copied to the host on the
+  /// current stream.
+  Results read_back(const bool source_csr = true,
+                    const bool virial = true) const {
+    const auto host_vector = [](const torch::Tensor& tensor) {
+      const torch::Tensor values = tensor.cpu();
+      return std::vector(values.data_ptr(),
+                                 values.data_ptr() + values.numel());
+    };
+    Results results;
+    results.atom_energy = host_vector(out_atom_energy);
+    results.force = host_vector(out_force);
+    if constexpr (Case::kSpin) {
+      results.force_mag = host_vector(out_force_mag);
+    }
+    if (virial) {
+      results.virial = host_vector(out_virial);
+    }
+    if (source_csr) {
+      results.atom_virial = host_vector(out_atom_virial);
+    }
+    return results;
+  }
+
+  /// Evaluate the graph on the current stream.
+  Results evaluate(const bool source_csr = true, const bool virial = true) {
+    reset_outputs();
+    run_device_entry(nullptr, source_csr, virial);
+    return read_back(source_csr, virial);
+  }
+#endif
+};
+
+TYPED_TEST_SUITE_P(TestInferDpa4cCanonicalPtExpt);
+
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt, host_entry_matches_reference) {
+  this->expect_reference(this->host, "default_state");
+}
+
+#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt, device_entry_matches_host_entry) {
+  // The first call runs eagerly, the second captures and replays, the third
+  // replays.
+  for (int call = 0; call < 3; ++call) {
+    this->expect_equal(this->evaluate(), this->host,
+                       "call " + std::to_string(call));
+  }
+}
+
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt,
+             a_caller_stream_orders_the_evaluation) {
+#ifdef DEEPMD_TORCH_HAS_CUDA
+  // The outputs are reset and read back on the caller's stream alone, after
+  // the stream is held back, while the current stream of the thread stays
+  // another one. A replay enqueued anywhere but on the caller's stream
+  // finishes before the reset and reads back as NaN.
+  const c10::cuda::CUDAStream stream =
+      c10::cuda::getStreamFromPool(/*isHighPriority=*/false, /*device=*/0);
+  for (int call = 0; call < 3; ++call) {
+    {
+      const c10::cuda::CUDAStreamGuard guard(stream);
+      at::cuda::sleep(kHoldCycles);
+      this->reset_outputs();
+    }
+    this->run_device_entry(stream.stream());
+    const c10::cuda::CUDAStreamGuard guard(stream);
+    this->expect_equal(this->read_back(), this->host,
+                       "call " + std::to_string(call) + " on a caller stream");
+  }
+#else
+  GTEST_SKIP() << "Skip because deepmd-kit was built without the CUDA stream "
+                  "API of PyTorch.";
+#endif
+}
+
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt, the_source_csr_is_optional) {
+  if constexpr (!TypeParam::kSpin) {
+    ASSERT_FALSE(this->dp.canonical_graph_requires_source_csr());
+  }
+  for (int call = 0; call < 3; ++call) {
+    this->expect_equal(
+        this->evaluate(/*source_csr=*/false), this->host,
+        "call " + std::to_string(call) + " without the source CSR");
+  }
+}
+
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt, the_virial_is_optional) {
+  for (int call = 0; call < 3; ++call) {
+    this->expect_equal(this->evaluate(/*source_csr=*/true, /*virial=*/false),
+                       this->host,
+                       "call " + std::to_string(call) + " without the virial");
+  }
+}
+
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt,
+             a_replay_reads_the_inputs_in_place) {
+  // The system is strained, as a barostat rescales it between two neighbor
+  // list rebuilds of an MD run, then the types of the atoms change, and for
+  // the spin model then the moments. Each configuration is rewritten into the
+  // graph in place, so that its evaluations replay the graph captured for
+  // the first one.
+  std::vector steps(1, this->system);
+  for (double& value : steps[0].coord) {
+    value *= 1.0 + kStrain;
+  }
+  for (double& value : steps[0].box) {
+    value *= 1.0 + kStrain;
+  }
+  steps.push_back(steps[0]);
+  for (int& type : steps[1].atype) {
+    type = 1 - type;
+  }
+  if constexpr (TypeParam::kSpin) {
+    steps.push_back(steps[1]);
+    for (double& moment : steps[2].spin) {
+      moment = -0.5 * moment;
+    }
+  }
+
+  // The first call runs eagerly, the second captures and replays.
+  for (int call = 0; call < 2; ++call) {
+    this->expect_equal(this->evaluate(), this->host,
+                       "call " + std::to_string(call));
+  }
+  Results previous = this->host;
+  for (std::size_t step = 0; step < steps.size(); ++step) {
+    const Results expected = this->compute_host(this->dp, steps[step]);
+    // Each change must alter the results, or a replay that read the inputs
+    // of the capture would still match.
+    const double change =
+        std::max({max_difference(expected.atom_energy, previous.atom_energy),
+                  max_difference(expected.force, previous.force),
+                  max_difference(expected.force_mag, previous.force_mag)});
+    ASSERT_GT(change, 1e-6) << "step " << step;
+    ASSERT_NO_FATAL_FAILURE(this->rewrite_in_place(steps[step]));
+    for (int call = 0; call < 2; ++call) {
+      this->expect_equal(this->evaluate(), expected,
+                         "replay " + std::to_string(call) + " after step " +
+                             std::to_string(step));
+    }
+    previous = expected;
+  }
+}
+
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt, a_new_key_releases_the_replay) {
+  // Outputs at other addresses form another key: the graph captured for the
+  // first outputs is released, and the evaluations under the new key run
+  // eagerly, are captured and are replayed in turn, as they are again once
+  // the first outputs return.
+  for (int call = 0; call < 2; ++call) {
+    this->expect_equal(this->evaluate(), this->host,
+                       "call " + std::to_string(call));
+  }
+  const std::vector first = {
+      this->out_atom_energy, this->out_force, this->out_force_mag,
+      this->out_virial, this->out_atom_virial};
+  this->allocate_outputs();
+  ASSERT_NE(this->out_force.data_ptr(), first[1].data_ptr());
+  for (int call = 0; call < 3; ++call) {
+    this->expect_equal(this->evaluate(), this->host,
+                       "call " + std::to_string(call) + " on other outputs");
+  }
+  this->out_atom_energy = first[0];
+  this->out_force = first[1];
+  this->out_force_mag = first[2];
+  this->out_virial = first[3];
+  this->out_atom_virial = first[4];
+  for (int call = 0; call < 3; ++call) {
+    this->expect_equal(
+        this->evaluate(), this->host,
+        "call " + std::to_string(call) + " back on the first outputs");
+  }
+}
+
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt,
+             a_graph_with_halo_nodes_matches_the_host_entry) {
+  // The unfolded graph gives every ghost atom a node, whose forces the caller
+  // adds to those of its owner, as the reverse communication of domain
+  // decomposition does; the extents of such a graph change with every
+  // neighbor rebuild, so it is evaluated eagerly. Its rows are those of the
+  // folded graph, so its energies, folded forces and virial agree with the
+  // host entry points up to the summation order of the run, which the node
+  // count of the graph configures.
+  const auto graph = this->build_graph(this->system, /*fold=*/false);
+  this->use_graph(graph);
+  ASSERT_GT(this->nodes, this->natoms);
+  this->allocate_outputs();
+  const Results& host = this->host;
+  for (int call = 0; call < 2; ++call) {
+    const Results halo = this->evaluate();
+    const std::string what = "call " + std::to_string(call);
+    EXPECT_LT(max_difference(halo.atom_energy, host.atom_energy),
+              kReferenceTolerance * largest(host.atom_energy))
+        << what << ": atom energy";
+    EXPECT_LT(
+        max_difference(fold_onto_owners(halo.force, graph.owner, this->natoms),
+                       host.force),
+        kReferenceTolerance * largest(host.force))
+        << what << ": force";
+    EXPECT_LT(max_difference(halo.virial, host.virial),
+              kReferenceTolerance * largest(host.virial))
+        << what << ": virial";
+    if constexpr (TypeParam::kSpin) {
+      EXPECT_LT(max_difference(
+                    fold_onto_owners(halo.force_mag, graph.owner, this->natoms),
+                    host.force_mag),
+                kReferenceTolerance * largest(host.force_mag))
+          << what << ": magnetic force";
+    }
+  }
+}
+
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt, a_capture_elsewhere_stays_valid) {
+#ifdef DEEPMD_TORCH_HAS_CUDA
+  // Another thread holds a capture of its own open while this thread replays
+  // the graph of one rank and evaluates a graph with halo nodes eagerly, both
+  // without a caller's stream. Neither waits for the whole device, which is
+  // invalid while a capture is under way and would invalidate that capture.
+  for (int call = 0; call < 2; ++call) {
+    this->expect_equal(this->evaluate(), this->host,
+                       "call " + std::to_string(call));
+  }
+  const auto halo = this->build_graph(this->system, /*fold=*/false);
+  torch::cuda::synchronize();
+  CaptureElsewhere capture;
+  this->expect_equal(this->evaluate(), this->host,
+                     "a replay while another thread captures");
+  this->use_graph(halo);
+  this->allocate_outputs();
+  const Results eager = this->evaluate();
+  EXPECT_EQ(capture.close(), cudaSuccess)
+      << "an evaluation invalidated the capture of another thread";
+  EXPECT_LT(
+      max_difference(fold_onto_owners(eager.force, halo.owner, this->natoms),
+                     this->host.force),
+      kReferenceTolerance * largest(this->host.force))
+      << "an eager evaluation while another thread captures";
+#else
+  GTEST_SKIP() << "Skip because deepmd-kit was built without the CUDA stream "
+                  "API of PyTorch.";
+#endif
+}
+
+TYPED_TEST_P(TestInferDpa4cCanonicalPtExpt, a_charge_state_reaches_the_replay) {
+  // A fresh instance evaluated through the host entry points is the
+  // reference under the other state.
+  typename TypeParam::Evaluator other;
+  other.init(TypeParam::kModelPath);
+  other.set_charge_spin(kOtherChargeSpin);
+  const Results other_host = this->compute_host(other, this->system);
+  ASSERT_GT(max_difference(other_host.atom_energy, this->host.atom_energy),
+            1e-6)
+      << "the two charge states must differ, or a replay that kept the "
+         "constants of the first state would still match";
+  this->expect_reference(other_host, "other_state");
+
+  for (int call = 0; call < 3; ++call) {
+    this->expect_equal(this->evaluate(), this->host,
+                       "call " + std::to_string(call) + " before the state");
+  }
+  this->dp.set_charge_spin(kOtherChargeSpin);
+  for (int call = 0; call < 3; ++call) {
+    this->expect_equal(this->evaluate(), other_host,
+                       "call " + std::to_string(call) + " after the state");
+  }
+}
+
+REGISTER_TYPED_TEST_SUITE_P(TestInferDpa4cCanonicalPtExpt,
+                            host_entry_matches_reference,
+                            device_entry_matches_host_entry,
+                            a_caller_stream_orders_the_evaluation,
+                            the_source_csr_is_optional,
+                            the_virial_is_optional,
+                            a_replay_reads_the_inputs_in_place,
+                            a_new_key_releases_the_replay,
+                            a_graph_with_halo_nodes_matches_the_host_entry,
+                            a_capture_elsewhere_stays_valid,
+                            a_charge_state_reaches_the_replay);
+#else
+REGISTER_TYPED_TEST_SUITE_P(TestInferDpa4cCanonicalPtExpt,
+                            host_entry_matches_reference);
+#endif
+
+}  // namespace dpa4c_canonical
+}  // namespace deepmd_test
diff --git a/source/api_cc/tests/test_deeppot_dpa4c_canonical_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4c_canonical_ptexpt.cc
new file mode 100644
index 0000000000..31ad21d775
--- /dev/null
+++ b/source/api_cc/tests/test_deeppot_dpa4c_canonical_ptexpt.cc
@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// The device-resident entry points of the compressed DPA4C archives frozen
+// with the compact canonical lower, through deepmd::DeepPot and
+// deepmd::DeepSpin (see dpa4c_canonical_test_common.h for the cases).
+#include 
+
+#include "DeepPot.h"
+#include "DeepSpin.h"
+#include "test_utils.h"
+// The shared cases follow the helpers of test_utils.h, which they call.
+#include "dpa4c_canonical_test_common.h"
+
+namespace deepmd_test {
+namespace dpa4c_canonical {
+
+using Models =
+    ::testing::Types, Spin>;
+INSTANTIATE_TYPED_TEST_SUITE_P(CXX,
+                               TestInferDpa4cCanonicalPtExpt,
+                               Models,
+                               CaseNames);
+
+}  // namespace dpa4c_canonical
+}  // namespace deepmd_test
diff --git a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc
new file mode 100644
index 0000000000..37b423a59a
--- /dev/null
+++ b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc
@@ -0,0 +1,907 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Test C++ inference for pt_expt (.pt2) DPA4C models with a latent-charge
+// long-range module, which sums over every pair of atoms of a complete
+// non-periodic system, and over the images of the cell of a periodic one when
+// the module serves periodic systems. Models generated by
+// source/tests/infer/gen_dpa4c_long_range.py: uncompressed archives on the
+// graph lower, which receives the pair list or the mesh of the cell, and
+// compressed archives on the compact canonical lower, which reads the
+// positions of the atoms and the cell.
+#include 
+
+#include 
+#include 
+#include 
+#include 
+
+#include "DeepPot.h"
+#include "DeepPotPTExpt.h"
+#include "expected_ref.h"
+#include "neighbor_list.h"
+#include "test_utils.h"
+
+#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT
+#include 
+
+#include 
+#include 
+#include 
+
+#include "../src/longRangePTExpt.h"
+#include "commonPT.h"
+#ifdef DEEPMD_TORCH_HAS_CUDA
+#include 
+#include 
+#include 
+#include 
+#endif
+#endif
+
+namespace {
+constexpr const char* kModelPath =
+    "../../tests/infer/deeppot_dpa4c_long_range.pt2";
+constexpr const char* kRefPath =
+    "../../tests/infer/deeppot_dpa4c_long_range.expected";
+constexpr const char* kCompressedModelPath =
+    "../../tests/infer/deeppot_dpa4c_long_range_compressed.pt2";
+constexpr const char* kCompressedRefPath =
+    "../../tests/infer/deeppot_dpa4c_long_range_compressed.expected";
+constexpr const char* kPeriodicModelPath =
+    "../../tests/infer/deeppot_dpa4c_long_range_periodic.pt2";
+constexpr const char* kPeriodicRefPath =
+    "../../tests/infer/deeppot_dpa4c_long_range_periodic.expected";
+constexpr const char* kPeriodicCompressedModelPath =
+    "../../tests/infer/deeppot_dpa4c_long_range_periodic_compressed.pt2";
+constexpr const char* kPeriodicCompressedRefPath =
+    "../../tests/infer/deeppot_dpa4c_long_range_periodic_compressed.expected";
+constexpr const char* kPeriodicMeshRefPath =
+    "../../tests/infer/deeppot_dpa4c_long_range_periodic_mesh.expected";
+
+// The full neighbor list of a non-periodic system within ``rc``: every atom
+// is local, so the list holds no ghost.
+template 
+std::vector> nopbc_nlist(const std::vector& coord,
+                                          const double rc) {
+  const int natoms = static_cast(coord.size() / 3);
+  std::vector> nlist(natoms);
+  for (int ii = 0; ii < natoms; ++ii) {
+    for (int jj = 0; jj < natoms; ++jj) {
+      double r2 = 0.;
+      for (int dd = 0; dd < 3; ++dd) {
+        const double diff = coord[ii * 3 + dd] - coord[jj * 3 + dd];
+        r2 += diff * diff;
+      }
+      if (jj != ii && r2 < rc * rc) {
+        nlist[ii].push_back(jj);
+      }
+    }
+  }
+  return nlist;
+}
+}  // namespace
+
+// One reference case of an archive: per-atom energies, forces and virials,
+// with the totals the interface returns.
+template 
+struct LongRangeRef {
+  std::vector expected_f;
+  double expected_tot_e = 0.;
+  std::vector expected_tot_v = std::vector(9, 0.);
+
+  void load(const deepmd_test::ExpectedRef& ref, const char* section) {
+    const std::vector atom_e =
+        ref.template get(section, "expected_e");
+    const std::vector atom_v =
+        ref.template get(section, "expected_v");
+    expected_f = ref.template get(section, "expected_f");
+    for (std::size_t ii = 0; ii < atom_e.size(); ++ii) {
+      expected_tot_e += atom_e[ii];
+      for (int dd = 0; dd < 9; ++dd) {
+        expected_tot_v[dd] += atom_v[ii * 9 + dd];
+      }
+    }
+  }
+
+  void check(double ener,
+             const std::vector& force,
+             const std::vector& virial,
+             double epsilon) const {
+    ASSERT_EQ(force.size(), expected_f.size());
+    ASSERT_EQ(virial.size(), 9u);
+    EXPECT_LT(fabs(ener - expected_tot_e), epsilon);
+    for (std::size_t ii = 0; ii < force.size(); ++ii) {
+      EXPECT_LT(fabs(force[ii] - expected_f[ii]), epsilon);
+    }
+    for (int ii = 0; ii < 9; ++ii) {
+      EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), epsilon);
+    }
+  }
+};
+
+template 
+class TestInferDeepPotDpa4cLongRangePtExpt : public ::testing::Test {
+ protected:
+  std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74,
+                                  0.25,  3.32, 1.68, 3.36,  3.00, 1.81,
+                                  3.51,  2.51, 2.60, 4.27,  3.22, 1.56};
+  std::vector atype = {0, 1, 1, 0, 1, 1};
+  std::vector charge_spin_explicit = {1.0, 2.0};
+  LongRangeRef ref_default;
+  LongRangeRef ref_explicit;
+
+  static deepmd::DeepPot dp;
+
+  static void SetUpTestSuite() {
+#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT
+    dp.init(kModelPath);
+#endif
+  }
+
+  void SetUp() override {
+#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT
+    GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#endif
+    deepmd_test::ExpectedRef ref;
+    ref.load(kRefPath);
+    ref_default.load(ref, "default");
+    ref_explicit.load(ref, "explicit");
+  }
+
+  static void TearDownTestSuite() { dp = deepmd::DeepPot(); }
+
+  // Computes through the LAMMPS-style entry point, with the full neighbor
+  // list of the system and no ghost atom.
+  void compute_lmp(double& ener,
+                   std::vector& force,
+                   std::vector& virial,
+                   const std::vector& charge_spin) {
+    std::vector> nlist_data = nopbc_nlist(coord, dp.cutoff());
+    const int nloc = static_cast(atype.size());
+    std::vector ilist(nloc), numneigh(nloc);
+    std::vector firstneigh(nloc);
+    deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]);
+    convert_nlist(inlist, nlist_data);
+    dp.compute(ener, force, virial, coord, atype, std::vector(), 0,
+               inlist, 0, {}, {}, charge_spin);
+  }
+};
+
+template 
+deepmd::DeepPot TestInferDeepPotDpa4cLongRangePtExpt::dp;
+
+TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangePtExpt, ValueTypes);
+
+TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, has_long_range) {
+  EXPECT_TRUE(this->dp.has_long_range());
+  EXPECT_EQ(this->dp.dim_chg_spin(), 2);
+}
+
+TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, standalone_default) {
+  using VALUETYPE = TypeParam;
+  const double epsilon = std::is_same::value ? 1e-8 : 1e-4;
+  double ener;
+  std::vector force, virial;
+  this->dp.compute(ener, force, virial, this->coord, this->atype,
+                   std::vector());
+  this->ref_default.check(ener, force, virial, epsilon);
+}
+
+TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, standalone_explicit) {
+  using VALUETYPE = TypeParam;
+  const double epsilon = std::is_same::value ? 1e-8 : 1e-4;
+  double ener;
+  std::vector force, virial;
+  this->dp.compute(ener, force, virial, this->coord, this->atype,
+                   std::vector(), {}, {},
+                   this->charge_spin_explicit);
+  this->ref_explicit.check(ener, force, virial, epsilon);
+}
+
+TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, lmp_nlist_explicit) {
+  using VALUETYPE = TypeParam;
+  const double epsilon = std::is_same::value ? 1e-8 : 1e-4;
+  double ener;
+  std::vector force, virial;
+  this->compute_lmp(ener, force, virial, this->charge_spin_explicit);
+  this->ref_explicit.check(ener, force, virial, epsilon);
+}
+
+// A module that serves no periodic system rejects a cell.
+TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, periodic_box_is_rejected) {
+  using VALUETYPE = TypeParam;
+  const std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.};
+  double ener;
+  std::vector force, virial;
+  EXPECT_THROW(
+      this->dp.compute(ener, force, virial, this->coord, this->atype, box),
+      deepmd::deepmd_exception);
+}
+
+// The compressed archive folds the charge state into constants, which
+// set_charge_spin rewrites: the tables of the descriptor and the total charge
+// of the constraint together.
+template 
+class TestInferDeepPotDpa4cLongRangeCompressedPtExpt : public ::testing::Test {
+ protected:
+  std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74,
+                                  0.25,  3.32, 1.68, 3.36,  3.00, 1.81,
+                                  3.51,  2.51, 2.60, 4.27,  3.22, 1.56};
+  std::vector atype = {0, 1, 1, 0, 1, 1};
+  LongRangeRef ref_default;
+  LongRangeRef ref_explicit;
+  deepmd::DeepPot dp;
+
+  void SetUp() override {
+#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT
+    GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#endif
+    std::ifstream model_file(kCompressedModelPath);
+    if (!model_file.good()) {
+      GTEST_SKIP() << "Skip because " << kCompressedModelPath
+                   << " was not generated (gen_dpa4c_long_range.py writes it "
+                      "on a CUDA device only).";
+    }
+    dp.init(kCompressedModelPath);
+    deepmd_test::ExpectedRef ref;
+    ref.load(kCompressedRefPath);
+    ref_default.load(ref, "default");
+    ref_explicit.load(ref, "explicit");
+  }
+
+  void compute_lmp(double& ener,
+                   std::vector& force,
+                   std::vector& virial) {
+    std::vector> nlist_data = nopbc_nlist(coord, dp.cutoff());
+    const int nloc = static_cast(atype.size());
+    std::vector ilist(nloc), numneigh(nloc);
+    std::vector firstneigh(nloc);
+    deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]);
+    convert_nlist(inlist, nlist_data);
+    dp.compute(ener, force, virial, coord, atype, std::vector(), 0,
+               inlist, 0);
+  }
+};
+
+TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangeCompressedPtExpt, ValueTypes);
+
+TYPED_TEST(TestInferDeepPotDpa4cLongRangeCompressedPtExpt,
+           set_charge_spin_follows_the_state) {
+  using VALUETYPE = TypeParam;
+  constexpr double epsilon = 1e-4;
+  EXPECT_TRUE(this->dp.has_long_range());
+  double ener;
+  std::vector force, virial;
+  this->compute_lmp(ener, force, virial);
+  this->ref_default.check(ener, force, virial, epsilon);
+  this->dp.set_charge_spin({1.0, 2.0});
+  this->compute_lmp(ener, force, virial);
+  this->ref_explicit.check(ener, force, virial, epsilon);
+  this->dp.set_charge_spin({0.0, 1.0});
+  this->compute_lmp(ener, force, virial);
+  this->ref_default.check(ener, force, virial, epsilon);
+}
+
+// A module that serves periodic systems evaluates a periodic system over the
+// images of its cell, through the coordinates and the cell or through the
+// neighbor list of a LAMMPS-style call, whose ghost atoms are images, and
+// keeps the direct sum of a system without a cell. ``Archive`` selects the
+// uncompressed archive on the graph lower or, on a CUDA device, the
+// compressed one on the compact canonical lower.
+template 
+class TestInferDeepPotDpa4cLongRangePeriodicPtExpt : public ::testing::Test {
+ protected:
+  std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74,
+                                  0.25,  3.32, 1.68, 3.36,  3.00, 1.81,
+                                  3.51,  2.51, 2.60, 4.27,  3.22, 1.56};
+  std::vector atype = {0, 1, 1, 0, 1, 1};
+  std::vector box = {13.0, 0.0,  0.0, 0.6, 12.5,
+                                0.0,  -0.4, 0.8, 13.5};
+  LongRangeRef ref_periodic;
+  LongRangeRef ref_cluster;
+  deepmd::DeepPot dp;
+
+  void load(const char* model_path, const char* ref_path) {
+#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT
+    GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#endif
+    std::ifstream model_file(model_path);
+    if (!model_file.good()) {
+      GTEST_SKIP() << "Skip because " << model_path
+                   << " was not generated (gen_dpa4c_long_range.py writes the "
+                      "compressed archives on a CUDA device only, and the "
+                      "graph lower with torch >= 2.12).";
+    }
+    dp.init(model_path);
+    deepmd_test::ExpectedRef ref;
+    ref.load(ref_path);
+    ref_periodic.load(ref, "periodic");
+    ref_cluster.load(ref, "cluster");
+  }
+
+  // Computes the periodic system through the LAMMPS-style entry point, with
+  // the ghost atoms of the cell, and folds their forces back onto their
+  // owners.
+  void compute_lmp_periodic(double& ener,
+                            std::vector& force,
+                            std::vector& virial) {
+    std::vector coord_cpy;
+    std::vector atype_cpy, mapping;
+    std::vector> nlist_data;
+    _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord,
+                            atype, box, static_cast(dp.cutoff()));
+    const int nloc = static_cast(atype.size());
+    const int nall = static_cast(atype_cpy.size());
+    std::vector ilist(nloc), numneigh(nloc);
+    std::vector firstneigh(nloc);
+    deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]);
+    convert_nlist(inlist, nlist_data);
+    inlist.mapping = mapping.data();
+    std::vector force_all;
+    dp.compute(ener, force_all, virial, coord_cpy, atype_cpy, box, nall - nloc,
+               inlist, 0);
+    _fold_back(force, force_all, mapping, nloc, nall, 3);
+  }
+
+  void check_all(const double epsilon) {
+    EXPECT_TRUE(dp.has_long_range());
+    double ener;
+    std::vector force, virial;
+    dp.compute(ener, force, virial, coord, atype, box);
+    ref_periodic.check(ener, force, virial, epsilon);
+    compute_lmp_periodic(ener, force, virial);
+    ref_periodic.check(ener, force, virial, epsilon);
+    dp.compute(ener, force, virial, coord, atype, std::vector());
+    ref_cluster.check(ener, force, virial, epsilon);
+  }
+};
+
+TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, ValueTypes);
+
+TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, graph_lower) {
+  using VALUETYPE = TypeParam;
+  this->load(kPeriodicModelPath, kPeriodicRefPath);
+  if (this->IsSkipped()) {
+    return;
+  }
+  this->check_all(std::is_same::value ? 1e-8 : 1e-4);
+}
+
+TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, compact_lower) {
+  this->load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath);
+  if (this->IsSkipped()) {
+    return;
+  }
+  this->check_all(1e-4);
+}
+
+// A cell without volume marks no system: it is rejected rather than taken as
+// a non-periodic one.
+TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt,
+           degenerate_cell_is_rejected) {
+  using VALUETYPE = TypeParam;
+  this->load(kPeriodicModelPath, kPeriodicRefPath);
+  if (this->IsSkipped()) {
+    return;
+  }
+  double ener;
+  std::vector force, virial;
+  EXPECT_THROW(this->dp.compute(ener, force, virial, this->coord, this->atype,
+                                std::vector(9, 0.)),
+               deepmd::deepmd_exception);
+}
+
+// The mesh topology that the C++ interface builds for the graph lower of a
+// periodic model matches the mesh of the Python builder field by field, in
+// the cell and without a cell. The reference needs no archive, so the test
+// runs whatever PyTorch release generated the fixtures.
+TEST(TestInferDeepPotDpa4cLongRangeMeshPtExpt, matches_python_builder) {
+#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT
+  GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#else
+  deepmd_test::ExpectedRef ref;
+  ref.load(kPeriodicMeshRefPath);
+  // The fields in the order of the topology inputs of the graph lower.
+  const std::vector names = {"box",  "volume", "coord",  "strain",
+                                          "mask", "mode_1", "mode_2", "mode_3"};
+  std::vector coord = ref.get("periodic", "coord");
+  const auto coord_tensor =
+      torch::from_blob(coord.data(),
+                       {static_cast(coord.size() / 3), 3},
+                       torch::kFloat64)
+          .clone();
+  const double spacing = ref.get("periodic", "spacing")[0];
+  for (const std::string section : {"periodic", "cluster"}) {
+    const std::vector box =
+        section == "periodic" ? ref.get(section, "box") : std::vector{};
+    const auto fields = deepmd::ptexpt::buildMeshTopology(
+        coord_tensor, box, spacing, torch::kFloat64);
+    ASSERT_EQ(fields.size(), names.size());
+    for (std::size_t k = 0; k < names.size(); ++k) {
+      const auto actual = fields[k].to(torch::kFloat64).contiguous().view(-1);
+      const std::vector expected = ref.get(section, names[k]);
+      ASSERT_EQ(actual.numel(), static_cast(expected.size()))
+          << section << " " << names[k];
+      const double* value = actual.data_ptr();
+      for (std::size_t i = 0; i < expected.size(); ++i) {
+        EXPECT_NEAR(value[i], expected[i],
+                    1e-12 * std::max(1.0, std::abs(expected[i])))
+            << section << " " << names[k] << "[" << i << "]";
+      }
+    }
+  }
+#endif
+}
+
+#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT && defined(DEEPMD_TORCH_HAS_CUDA)
+// The device-resident entry point of the compressed archives, which reads the
+// positions of the nodes and the cell of a periodic system besides the compact
+// canonical graph. A graph without halo nodes is evaluated eagerly, captured
+// into a CUDA graph and replayed: that of a non-periodic system that one rank
+// holds whole, and that of a periodic system on the local atoms with the ghost
+// atoms folded onto their owners, as one rank builds it, whose replay evaluates
+// the cell of its call while the mesh of the cell keeps its size. A periodic
+// graph with a node for every ghost atom, an image of its owner, is evaluated
+// eagerly. Every evaluation writes into outputs reset to NaN, so an unwritten
+// output fails the comparison, and a replay allocates no device memory, which
+// tells it from an eager evaluation or a capture.
+class TestInferDeepPotDpa4cLongRangeDevicePtExpt : public ::testing::Test {
+ protected:
+  // A replay whose capture fails evaluates eagerly instead, with the same
+  // results; only this warning tells the two apart.
+  static constexpr const char* kCaptureWarning =
+      "could not capture the evaluation";
+  // Clock cycles for which at::cuda::sleep holds a stream back: tens of
+  // milliseconds, far longer than enqueueing an evaluation takes.
+  static constexpr std::int64_t kHoldCycles = std::int64_t{1} << 27;
+
+  std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74,
+                               0.25,  3.32, 1.68, 3.36,  3.00, 1.81,
+                               3.51,  2.51, 2.60, 4.27,  3.22, 1.56};
+  std::vector atype = {0, 1, 1, 0, 1, 1};
+  std::vector box = {13.0, 0.0, 0.0, 0.6, 12.5, 0.0, -0.4, 0.8, 13.5};
+  deepmd::DeepPot dp;
+  deepmd_test::ExpectedRef ref;
+  bool capturing_stderr = false;
+
+  // A compact canonical graph on the device with the positions of its nodes,
+  // the atom that owns every node, and outputs reused by its evaluations,
+  // whose addresses the replay of repeated evaluations keys on.
+  struct Graph {
+    deepmd::CanonicalGraphTensorPack tensors;
+    torch::Tensor coord;
+    std::vector owner;
+    int nodes = 0;
+    torch::Tensor atom_energy;
+    torch::Tensor force;
+    torch::Tensor virial;
+    torch::Tensor atom_virial;
+  };
+
+  // The outputs of an evaluation on the host, the forces folded onto the
+  // owners of the nodes, and the number of blocks that the caching allocator
+  // handed out on the device during the evaluation.
+  struct Results {
+    std::vector atom_energy;
+    std::vector force;
+    std::vector virial;
+    std::int64_t allocations = 0;
+  };
+
+  static std::int64_t device_allocations() {
+    const auto stats = c10::cuda::CUDACachingAllocator::getDeviceStats(0);
+    return stats
+        .allocation[static_cast(
+            c10::CachingAllocator::StatType::AGGREGATE)]
+        .allocated;
+  }
+
+  void load(const char* model_path, const char* ref_path) {
+    if (!torch::cuda::is_available()) {
+      GTEST_SKIP() << "Skip because the compact canonical path needs a CUDA "
+                      "device.";
+    }
+    std::ifstream model_file(model_path);
+    if (!model_file.good()) {
+      GTEST_SKIP() << "Skip because " << model_path
+                   << " was not generated (gen_dpa4c_long_range.py writes the "
+                      "compressed archives on a CUDA device only).";
+    }
+    dp.init(model_path);
+    ASSERT_TRUE(dp.uses_canonical_graph_inference());
+    ASSERT_TRUE(dp.has_long_range());
+    // The fused long-range operator serves the charge network of one channel
+    // and reads the source CSR for the per-atom virial alone.
+    ASSERT_FALSE(dp.canonical_graph_requires_source_csr());
+    ref.load(ref_path);
+    testing::internal::CaptureStderr();
+    capturing_stderr = true;
+  }
+
+  void TearDown() override {
+    if (!capturing_stderr) {
+      return;
+    }
+    const std::string log = testing::internal::GetCapturedStderr();
+    std::cerr << log;
+    EXPECT_EQ(log.find(kCaptureWarning), std::string::npos)
+        << "an evaluation fell back from the replay to eager evaluation";
+  }
+
+  // The graph of the neighbor list ``nlist`` of the extended atoms
+  // ``coord_ext``, of which the first ``nloc`` are local: folded onto the
+  // local atoms, or with a node for every extended atom.
+  Graph build_graph(const std::vector>& nlist,
+                    const std::vector& coord_ext,
+                    const std::vector& atype_ext,
+                    const std::vector& mapping,
+                    const int nloc,
+                    const bool fold) const {
+    const torch::Device device(torch::kCUDA, 0);
+    const int nall = static_cast(atype_ext.size());
+    const std::vector mapping_64(mapping.begin(), mapping.end());
+    deepmd::GraphTensorPack pack = deepmd::buildGraphTensors(
+        nlist, coord_ext, atype_ext, mapping_64, nloc, nall, dp.cutoff(),
+        device, /*row_centers=*/nullptr, fold);
+    deepmd::canonicalizeGraphPayload(pack, pack.atype.size(0));
+    pack.edge_vec = pack.edge_vec.to(torch::kFloat32);
+    Graph graph;
+    graph.tensors = deepmd::compactCanonicalGraph(pack);
+    deepmd::CanonicalGraphTensorPack& compact = graph.tensors;
+    compact.atype = compact.atype.to(torch::kInt64).contiguous();
+    compact.source = compact.source.contiguous();
+    compact.edge_vec = compact.edge_vec.contiguous();
+    compact.destination_row_ptr = compact.destination_row_ptr.contiguous();
+    compact.source_row_ptr = compact.source_row_ptr.contiguous();
+    compact.source_order = compact.source_order.contiguous();
+    graph.nodes = static_cast(compact.atype.size(0));
+    // The nodes are the leading extended atoms, in their order.
+    std::vector positions(
+        coord_ext.begin(), coord_ext.begin() + 3 * std::size_t(graph.nodes));
+    const auto f64 = torch::TensorOptions().dtype(torch::kFloat64);
+    graph.coord = torch::from_blob(positions.data(), {graph.nodes, 3}, f64)
+                      .to(device)
+                      .contiguous();
+    graph.owner = mapping;
+    graph.owner.resize(graph.nodes);
+    const auto f64_device = f64.device(device);
+    graph.atom_energy = torch::empty({nloc}, f64_device);
+    graph.force = torch::empty({graph.nodes, 3}, f64_device);
+    graph.virial = torch::empty({9}, f64_device);
+    graph.atom_virial = torch::empty({graph.nodes, 9}, f64_device);
+    return graph;
+  }
+
+  // ``graph`` with outputs of its own.
+  static Graph with_own_outputs(const Graph& graph) {
+    Graph copy = graph;
+    copy.atom_energy = torch::empty_like(graph.atom_energy);
+    copy.force = torch::empty_like(graph.force);
+    copy.virial = torch::empty_like(graph.virial);
+    copy.atom_virial = torch::empty_like(graph.atom_virial);
+    return copy;
+  }
+
+  // The cell scaled by ``factor``.
+  std::vector scaled_box(const double factor) const {
+    std::vector cell = box;
+    for (double& entry : cell) {
+      entry *= factor;
+    }
+    return cell;
+  }
+
+  // Enqueues on the current stream an evaluation of ``graph`` through
+  // compute_canonical_graph_gpu, in the cell ``cell`` (null for a
+  // non-periodic system), into outputs reset to NaN; ``stream`` is that
+  // stream, or null, with which the call waits for it before returning.
+  // Without ``source_csr`` the call passes neither the source CSR nor the
+  // per-atom virial.
+  void enqueue(Graph& graph,
+               const int nloc,
+               const double* cell,
+               void* stream,
+               const bool source_csr = true) {
+    const double nan = std::numeric_limits::quiet_NaN();
+    graph.atom_energy.fill_(nan);
+    graph.force.fill_(nan);
+    graph.virial.fill_(nan);
+    graph.atom_virial.fill_(nan);
+    const deepmd::CanonicalGraphTensorPack& compact = graph.tensors;
+    dp.compute_canonical_graph_gpu(
+        graph.atom_energy.data_ptr(), graph.force.data_ptr(),
+        graph.virial.data_ptr(),
+        source_csr ? graph.atom_virial.data_ptr() : nullptr,
+        compact.atype.data_ptr(), compact.source.data_ptr(),
+        deepmd::canonical_index_width(graph.nodes),
+        compact.edge_vec.data_ptr(),
+        compact.destination_row_ptr.data_ptr(),
+        source_csr ? compact.source_row_ptr.data_ptr() : nullptr,
+        source_csr ? compact.source_order.data_ptr() : nullptr,
+        graph.coord.data_ptr(), /*world=*/nullptr, cell, nloc,
+        graph.nodes, compact.source.size(0), stream);
+  }
+
+  // The atomic energies, forces and virial of an evaluation of ``graph`` on
+  // the host, the forces folded onto the owners of the nodes.
+  static Results collect(const Graph& graph,
+                         const int nloc,
+                         const torch::Tensor& atom_energy,
+                         const torch::Tensor& force,
+                         const torch::Tensor& virial) {
+    const auto host_vector = [](const torch::Tensor& tensor) {
+      const torch::Tensor values = tensor.cpu().contiguous();
+      return std::vector(values.data_ptr(),
+                                 values.data_ptr() + values.numel());
+    };
+    Results results;
+    results.atom_energy = host_vector(atom_energy);
+    results.virial = host_vector(virial);
+    const std::vector node_force = host_vector(force);
+    results.force.assign(3 * std::size_t(nloc), 0.0);
+    for (int node = 0; node < graph.nodes; ++node) {
+      for (int dd = 0; dd < 3; ++dd) {
+        results.force[3 * graph.owner[node] + dd] += node_force[3 * node + dd];
+      }
+    }
+    return results;
+  }
+
+  // Evaluates ``graph`` on the current stream, in the cell ``cell`` (null for
+  // a non-periodic system), with or without the source CSR.
+  Results evaluate(Graph& graph,
+                   const int nloc,
+                   const double* cell,
+                   const bool source_csr = true) {
+    const std::int64_t allocations = device_allocations();
+    enqueue(graph, nloc, cell, /*stream=*/nullptr, source_csr);
+    const std::int64_t allocated = device_allocations() - allocations;
+    Results results =
+        collect(graph, nloc, graph.atom_energy, graph.force, graph.virial);
+    results.allocations = allocated;
+    return results;
+  }
+
+  // Agreement of an evaluation with the reference of a section.
+  void expect_reference(const Results& results, const char* section) const {
+    LongRangeRef expected;
+    expected.load(ref, section);
+    double energy = 0.0;
+    for (const double value : results.atom_energy) {
+      energy += value;
+    }
+    expected.check(energy, results.force, results.virial, 1e-4);
+  }
+
+  // Evaluates the folded graph of the system, periodic in its cell or not,
+  // with the source CSR and then without it, which declines the per-atom
+  // virial alone: the energies, forces and virial agree bit for bit, and the
+  // evaluations without it are captured and replayed.
+  void expect_the_source_csr_optional(const bool periodic) {
+    const int nloc = static_cast(atype.size());
+    std::vector> nlist;
+    std::vector coord_ext = coord;
+    std::vector atype_ext = atype;
+    std::vector mapping(nloc);
+    for (int ii = 0; ii < nloc; ++ii) {
+      mapping[ii] = ii;
+    }
+    if (periodic) {
+      _build_nlist(nlist, coord_ext, atype_ext, mapping, coord, atype,
+                           box, static_cast(dp.cutoff()));
+    } else {
+      nlist = nopbc_nlist(coord, dp.cutoff());
+    }
+    Graph graph =
+        build_graph(nlist, coord_ext, atype_ext, mapping, nloc, /*fold=*/true);
+    const double* cell = periodic ? box.data() : nullptr;
+    const Results with_csr = evaluate(graph, nloc, cell);
+    expect_reference(with_csr, periodic ? "periodic" : "default");
+    for (int call = 0; call < 3; ++call) {
+      const Results without = evaluate(graph, nloc, cell, /*source_csr=*/false);
+      const std::string what =
+          "call " + std::to_string(call) + " without the source CSR";
+      EXPECT_EQ(without.atom_energy, with_csr.atom_energy) << what;
+      EXPECT_EQ(without.force, with_csr.force) << what;
+      EXPECT_EQ(without.virial, with_csr.virial) << what;
+      EXPECT_EQ(without.allocations == 0, call == 2)
+          << what << ": " << without.allocations << " device allocations";
+    }
+  }
+};
+
+TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, cluster_is_replayed) {
+  ASSERT_NO_FATAL_FAILURE(load(kCompressedModelPath, kCompressedRefPath));
+  if (IsSkipped()) {
+    return;
+  }
+  const int nloc = static_cast(atype.size());
+  std::vector mapping(nloc);
+  for (int ii = 0; ii < nloc; ++ii) {
+    mapping[ii] = ii;
+  }
+  Graph graph = build_graph(nopbc_nlist(coord, dp.cutoff()), coord, atype,
+                            mapping, nloc, /*fold=*/true);
+  ASSERT_EQ(graph.nodes, nloc);
+  // Eager, captured, and replayed from the third call on.
+  const Results eager = evaluate(graph, nloc, nullptr);
+  expect_reference(eager, "default");
+  EXPECT_GT(eager.allocations, 0);
+  for (int call = 1; call < 4; ++call) {
+    const Results again = evaluate(graph, nloc, nullptr);
+    EXPECT_EQ(again.atom_energy, eager.atom_energy) << "call " << call;
+    EXPECT_EQ(again.force, eager.force) << "call " << call;
+    EXPECT_EQ(again.virial, eager.virial) << "call " << call;
+    EXPECT_EQ(again.allocations == 0, call > 1)
+        << "call " << call << ": " << again.allocations
+        << " device allocations";
+  }
+  // The positions are part of the key of the replay: the same values at
+  // another address are evaluated anew.
+  graph.coord = graph.coord.clone();
+  const Results moved = evaluate(graph, nloc, nullptr);
+  EXPECT_EQ(moved.atom_energy, eager.atom_energy);
+  EXPECT_EQ(moved.force, eager.force);
+  EXPECT_GT(moved.allocations, 0);
+}
+
+TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, periodic_cell) {
+  ASSERT_NO_FATAL_FAILURE(
+      load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath));
+  if (IsSkipped()) {
+    return;
+  }
+  std::vector> nlist;
+  std::vector coord_ext;
+  std::vector atype_ext, mapping;
+  _build_nlist(nlist, coord_ext, atype_ext, mapping, coord, atype, box,
+                       static_cast(dp.cutoff()));
+  const int nloc = static_cast(atype.size());
+  // At the mesh spacing of the model, the mesh of the cell has 40 points
+  // along every lattice vector, as it has for the cell strained by 1%, while
+  // the cell stretched by 10% has a mesh of 42, 42 and 48 points.
+  const std::vector strained = scaled_box(1.01);
+  const std::vector stretched = scaled_box(1.1);
+  const auto expect_equal = [](const Results& actual, const Results& expected,
+                               const bool replayed, const std::string& what) {
+    EXPECT_EQ(actual.atom_energy, expected.atom_energy) << what;
+    EXPECT_EQ(actual.force, expected.force) << what;
+    EXPECT_EQ(actual.virial, expected.virial) << what;
+    EXPECT_EQ(actual.allocations == 0, replayed)
+        << what << ": " << actual.allocations << " device allocations";
+  };
+
+  // The folded graph: the strained cell is evaluated eagerly, and the first
+  // cell, under the same key, is captured and replayed. The strained cell is
+  // then replayed at once, since its mesh keeps the size of the capture.
+  Graph graph = build_graph(nlist, coord_ext, atype_ext, mapping, nloc, true);
+  ASSERT_EQ(graph.nodes, nloc);
+  const Results strained_eager = evaluate(graph, nloc, strained.data());
+  EXPECT_GT(strained_eager.allocations, 0);
+  const Results first = evaluate(graph, nloc, box.data());
+  expect_reference(first, "periodic");
+  EXPECT_GT(first.allocations, 0);
+  EXPECT_NE(first.atom_energy, strained_eager.atom_energy);
+  expect_equal(evaluate(graph, nloc, box.data()), first, true, "replay");
+  expect_equal(evaluate(graph, nloc, strained.data()), strained_eager, true,
+               "replay in the strained cell");
+  expect_equal(evaluate(graph, nloc, box.data()), first, true,
+               "replay back in the first cell");
+
+  // A mesh of another size is evaluated eagerly until it has held for 64
+  // calls, then once more under its new key, captured at the 66th call and
+  // replayed from the 67th on.
+  const Results stretched_eager = evaluate(graph, nloc, stretched.data());
+  EXPECT_GT(stretched_eager.allocations, 0);
+  EXPECT_NE(stretched_eager.atom_energy, first.atom_energy);
+  for (int call = 2; call <= 67; ++call) {
+    expect_equal(evaluate(graph, nloc, stretched.data()), stretched_eager,
+                 call == 67,
+                 "call " + std::to_string(call) + " in the stretched cell");
+  }
+
+  // The unfolded graph has halo nodes and is evaluated eagerly.
+  Graph unfolded =
+      build_graph(nlist, coord_ext, atype_ext, mapping, nloc, false);
+  ASSERT_EQ(unfolded.nodes, static_cast(atype_ext.size()));
+  const Results unfolded_first = evaluate(unfolded, nloc, box.data());
+  expect_reference(unfolded_first, "periodic");
+  for (int call = 1; call < 3; ++call) {
+    expect_equal(evaluate(unfolded, nloc, box.data()), unfolded_first, false,
+                 "call " + std::to_string(call) + " of the unfolded graph");
+  }
+}
+
+TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt,
+       a_new_cell_waits_for_the_evaluations) {
+  ASSERT_NO_FATAL_FAILURE(
+      load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath));
+  if (IsSkipped()) {
+    return;
+  }
+  std::vector> nlist;
+  std::vector coord_ext;
+  std::vector atype_ext, mapping;
+  _build_nlist(nlist, coord_ext, atype_ext, mapping, coord, atype, box,
+                       static_cast(dp.cutoff()));
+  const int nloc = static_cast(atype.size());
+  // The strained cell keeps the mesh of the cell.
+  const std::vector strained = scaled_box(1.01);
+  // Three graphs over the same inputs with outputs of their own: every call
+  // changes the key of the replay, so that every evaluation runs eagerly and
+  // no capture or release of a graph waits for the device.
+  Graph held_graph =
+      build_graph(nlist, coord_ext, atype_ext, mapping, nloc, /*fold=*/true);
+  Graph same_graph = with_own_outputs(held_graph);
+  Graph strained_graph = with_own_outputs(held_graph);
+  const c10::cuda::CUDAStream held =
+      c10::cuda::getStreamFromPool(/*isHighPriority=*/false, /*device=*/0);
+  const c10::cuda::CUDAStream other =
+      c10::cuda::getStreamFromPool(/*isHighPriority=*/false, /*device=*/0);
+  ASSERT_NE(held.stream(), other.stream());
+  // The cell on a held stream, the same cell on another stream, then the
+  // strained cell on that stream, each evaluation followed by copies of its
+  // outputs on its stream.
+  const auto run = [&](const bool hold) {
+    std::vector copies;
+    const auto run_on = [&](Graph& graph, const std::vector& cell,
+                            const c10::cuda::CUDAStream& stream) {
+      const c10::cuda::CUDAStreamGuard guard(stream);
+      enqueue(graph, nloc, cell.data(), stream.stream());
+      for (const torch::Tensor& output :
+           {graph.atom_energy, graph.force, graph.virial}) {
+        copies.push_back(output.clone());
+      }
+    };
+    if (hold) {
+      const c10::cuda::CUDAStreamGuard guard(held);
+      at::cuda::sleep(kHoldCycles);
+    }
+    run_on(held_graph, box, held);
+    run_on(same_graph, box, other);
+    run_on(strained_graph, strained, other);
+    torch::cuda::synchronize();
+    return copies;
+  };
+  // The first round runs the one-time work of the program. In the second,
+  // the held stream runs its evaluation after the strained cell has been
+  // handed over, and that call writes the cell only once the evaluations
+  // before it, the held one included, have read it.
+  run(/*hold=*/false);
+  const std::vector copies = run(/*hold=*/true);
+  const Results first = evaluate(held_graph, nloc, box.data());
+  expect_reference(first, "periodic");
+  const Results first_strained = evaluate(held_graph, nloc, strained.data());
+  const auto expect_equal = [&](const Graph& graph, const std::size_t at,
+                                const Results& expected,
+                                const std::string& what) {
+    const Results actual =
+        collect(graph, nloc, copies[at], copies[at + 1], copies[at + 2]);
+    EXPECT_EQ(actual.atom_energy, expected.atom_energy) << what;
+    EXPECT_EQ(actual.force, expected.force) << what;
+    EXPECT_EQ(actual.virial, expected.virial) << what;
+  };
+  expect_equal(held_graph, 0, first, "the cell on the held stream");
+  expect_equal(same_graph, 3, first, "the cell on the other stream");
+  expect_equal(strained_graph, 6, first_strained,
+               "the strained cell on the other stream");
+}
+
+TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt,
+       the_source_csr_is_optional_in_a_cluster) {
+  ASSERT_NO_FATAL_FAILURE(load(kCompressedModelPath, kCompressedRefPath));
+  if (IsSkipped()) {
+    return;
+  }
+  expect_the_source_csr_optional(/*periodic=*/false);
+}
+
+TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt,
+       the_source_csr_is_optional_in_a_cell) {
+  ASSERT_NO_FATAL_FAILURE(
+      load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath));
+  if (IsSkipped()) {
+    return;
+  }
+  expect_the_source_csr_optional(/*periodic=*/true);
+}
+#endif
diff --git a/source/api_cc/tests/test_deeppot_model_devi_ptexpt.cc b/source/api_cc/tests/test_deeppot_model_devi_ptexpt.cc
index 19fbcd0b99..67da88f177 100644
--- a/source/api_cc/tests/test_deeppot_model_devi_ptexpt.cc
+++ b/source/api_cc/tests/test_deeppot_model_devi_ptexpt.cc
@@ -89,6 +89,8 @@ TYPED_TEST(TestInferDeepPotModeDeviPtExpt, attrs) {
   EXPECT_EQ(dp_md.dim_fparam(), 1);
   EXPECT_EQ(dp_md.dim_aparam(), 1);
   EXPECT_TRUE(dp_md.has_default_fparam());
+  EXPECT_FALSE(dp0.has_long_range());
+  EXPECT_FALSE(dp_md.has_long_range());
 }
 
 TYPED_TEST(TestInferDeepPotModeDeviPtExpt, cpu_build_nlist) {
diff --git a/source/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc
index fa68b99fc2..b7e354d252 100644
--- a/source/api_cc/tests/test_neighbor_list_data.cc
+++ b/source/api_cc/tests/test_neighbor_list_data.cc
@@ -140,10 +140,6 @@ TEST(TestNeighborListData, RoundTripWithEmptyRows) {
 
 #ifdef BUILD_PYTORCH
 TEST(TestNeighborListData, CompactCanonicalGraphDropsMaskedGuards) {
-#if TORCH_VERSION_MAJOR < 2 || \
-    (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR < 3)
-  GTEST_SKIP() << "uint32 tensors require PyTorch 2.3 or later";
-#endif
   GraphTensorPack graph;
   graph.atype = torch::tensor({0}, torch::kInt64);
   graph.n_node = torch::tensor({1}, torch::kInt64);
@@ -158,13 +154,12 @@ TEST(TestNeighborListData, CompactCanonicalGraphDropsMaskedGuards) {
   graph.source_row_ptr = torch::tensor({0, 1}, torch::kInt64);
 
   const auto compact = compactCanonicalGraph(graph);
-  EXPECT_EQ(compact.source.scalar_type(), deepmd::canonicalGraphIndexType());
+  EXPECT_EQ(compact.source.scalar_type(), torch::kInt32);
   EXPECT_EQ(compact.source.numel(), 2);
   EXPECT_EQ(compact.edge_vec.scalar_type(), torch::kFloat32);
   EXPECT_EQ(compact.edge_vec.size(0), 2);
   EXPECT_TRUE(
-      torch::equal(compact.source_order,
-                   torch::tensor({0, 1}, deepmd::canonicalGraphIndexType())));
+      torch::equal(compact.source_order, torch::tensor({0, 1}, torch::kInt64)));
   EXPECT_EQ(compact.destination_row_ptr.select(0, 1).item(), 1);
   EXPECT_EQ(compact.source_row_ptr.select(0, 1).item(), 1);
 }
diff --git a/source/api_cc/tests/test_with_comm_load_failure_ptexpt.cc b/source/api_cc/tests/test_with_comm_load_failure_ptexpt.cc
index 768c1947bf..56b2b33b81 100644
--- a/source/api_cc/tests/test_with_comm_load_failure_ptexpt.cc
+++ b/source/api_cc/tests/test_with_comm_load_failure_ptexpt.cc
@@ -100,10 +100,11 @@ TEST_F(TestDeepPotPTExptWithCommLoadFailure, single_rank_compute_succeeds) {
 }
 
 TEST_F(TestDeepPotPTExptWithCommLoadFailure, multi_rank_compute_throws) {
-  // nswap > 0 forces the dispatch site to ``run_model_with_comm``; the
-  // load-failure guard added by PR #5430 must throw rather than silently
-  // falling back to the single-rank path. The send/recv arrays remain
-  // null — the guard fires before any of them are dereferenced.
+  // The multi-rank dispatch keys on nprocs (DeepPotPTExpt.cc): nprocs > 1
+  // forces the dispatch site to ``run_model_with_comm``, whose load-failure
+  // guard must throw rather than silently falling back to the single-rank
+  // path. The send/recv arrays remain null — the guard fires before any of
+  // them are dereferenced.
   float rc = dp.cutoff();
   int nloc = coord.size() / 3;
   std::vector coord_cpy;
@@ -118,7 +119,7 @@ TEST_F(TestDeepPotPTExptWithCommLoadFailure, multi_rank_compute_throws) {
                             firstneigh.data());
   convert_nlist(inlist, nlist_data);
   inlist.mapping = mapping.data();
-  inlist.nswap = 1;  // simulate multi-rank without populating send/recv
+  inlist.nprocs = 2;  // simulate multi-rank without populating send/recv
 
   double ener;
   std::vector force_, virial;
diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh
index 17b4f2f593..b8fe41b5ae 100755
--- a/source/install/test_cc_local.sh
+++ b/source/install/test_cc_local.sh
@@ -1,6 +1,27 @@
 #!/bin/bash
 set -ex
 
+# With no arguments, build, generate every fixture, and run CTest as before.
+# The independent Kokkos job needs the installation and only long-range models.
+MODE=${1:-all}
+if [ "$#" -gt 1 ]; then
+	echo "Usage: $0 [--build-only | --models-only=kokkos]" >&2
+	exit 2
+fi
+case "$MODE" in
+all | --build-only | --models-only=kokkos) ;;
+*)
+	echo "Usage: $0 [--build-only | --models-only=kokkos]" >&2
+	exit 2
+	;;
+esac
+if [ "$MODE" = "--models-only=kokkos" ]; then
+	if [ "${DP_VARIANT:-}" != "cuda" ] || [ "${ENABLE_PYTORCH:-TRUE}" != "TRUE" ]; then
+		echo "Kokkos fixtures require DP_VARIANT=cuda and ENABLE_PYTORCH=TRUE." >&2
+		exit 1
+	fi
+fi
+
 if [ "$DP_VARIANT" = "cuda" ]; then
 	CUDA_ARGS="-DUSE_CUDA_TOOLKIT=TRUE"
 elif [ "$DP_VARIANT" = "rocm" ]; then
@@ -20,19 +41,28 @@ BUILD_TMP_DIR=${SCRIPT_PATH}/../build_tests
 PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir
 mkdir -p ${BUILD_TMP_DIR}
 cd ${BUILD_TMP_DIR}
-cmake \
-	-D ENABLE_TENSORFLOW=${ENABLE_TENSORFLOW:-TRUE} \
-	-D ENABLE_PYTORCH=${ENABLE_PYTORCH:-TRUE} \
-	-D ENABLE_PADDLE=${ENABLE_PADDLE:-TRUE} \
-	-D INSTALL_TENSORFLOW=FALSE \
-	-D USE_TF_PYTHON_LIBS=${ENABLE_TENSORFLOW:-TRUE} \
-	-D USE_PT_PYTHON_LIBS=${ENABLE_PYTORCH:-TRUE} \
-	-D CMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} \
-	-D BUILD_TESTING:BOOL=TRUE \
-	-D LAMMPS_VERSION=stable_22Jul2025_update2 \
-	${CUDA_ARGS} ..
-cmake --build . -j${NPROC}
-cmake --install .
+if [ "$MODE" != "--models-only=kokkos" ]; then
+	BUILD_TESTING=TRUE
+	if [ "$MODE" = "--build-only" ]; then
+		BUILD_TESTING=FALSE
+	fi
+	cmake \
+		-D ENABLE_TENSORFLOW=${ENABLE_TENSORFLOW:-TRUE} \
+		-D ENABLE_PYTORCH=${ENABLE_PYTORCH:-TRUE} \
+		-D ENABLE_PADDLE=${ENABLE_PADDLE:-TRUE} \
+		-D INSTALL_TENSORFLOW=FALSE \
+		-D USE_TF_PYTHON_LIBS=${ENABLE_TENSORFLOW:-TRUE} \
+		-D USE_PT_PYTHON_LIBS=${ENABLE_PYTORCH:-TRUE} \
+		-D CMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} \
+		-D BUILD_TESTING:BOOL=${BUILD_TESTING} \
+		-D LAMMPS_VERSION=stable_22Jul2025_update2 \
+		${CUDA_ARGS} ..
+	cmake --build . -j${NPROC}
+	cmake --install .
+else
+	# Never let a missing installation fall back to an unrelated Python build.
+	test -f "${INSTALL_PREFIX}/lib/libdeepmd_c.so"
+fi
 # Generate PT/PT2 model files for C++ tests.
 # Must run after cmake --build so that libdeepmd_op_pt.so (custom ops) is available.
 if [ "${ENABLE_PYTORCH:-TRUE}" == "TRUE" ]; then
@@ -40,6 +70,9 @@ if [ "${ENABLE_PYTORCH:-TRUE}" == "TRUE" ]; then
 	# loads it via cxx_op.py.  The .so depends on libdeepmd.so (compute
 	# kernels) in the install prefix, so add that to LD_LIBRARY_PATH too.
 	export LD_LIBRARY_PATH=${INSTALL_PREFIX}/lib:${LD_LIBRARY_PATH}
+	if [ "$MODE" != "all" ]; then
+		test -f "${BUILD_TMP_DIR}/op/pt/libdeepmd_op_pt.so"
+	fi
 	python -c '
 import shutil, sys
 from pathlib import Path
@@ -55,13 +88,15 @@ else:
     shutil.copy2(str(so), str(dst))
     print(f"Installed {so} -> {dst}")
 '
+fi
+if [ "$MODE" = "--build-only" ]; then
+	exit 0
+fi
+if [ "${ENABLE_PYTORCH:-TRUE}" == "TRUE" ]; then
 	# When the build uses -fsanitize=leak, the custom op .so requires the LSAN
 	# runtime to be preloaded (otherwise dlopen fails).  We disable leak detection
 	# in the gen scripts to avoid false reports from torch/paddle internals.
 	INFER_SCRIPT_PATH=${SCRIPT_PATH}/../tests/infer
-	# Remove stale generated model files so they can't be accidentally reused
-	# if gen scripts change format or the code version changes.
-	rm -f ${INFER_SCRIPT_PATH}/*.pt2 ${INFER_SCRIPT_PATH}/*.pte
 	_GEN_ENV=""
 	if echo "${CXXFLAGS:-}" | grep -q fsanitize=leak; then
 		_LSAN_LIB=$(gcc -print-file-name=liblsan.so 2>/dev/null || true)
@@ -75,6 +110,17 @@ else:
 			_GEN_ENV="LD_PRELOAD=${_LSAN_LIB} LSAN_OPTIONS=detect_leaks=0 DP_GEN_UNDER_SANITIZER=lsan"
 		fi
 	fi
+	if [ "$MODE" = "--models-only=kokkos" ]; then
+		# Leave the other C++ fixtures alone when reusing an existing build.
+		rm -f ${INFER_SCRIPT_PATH}/deeppot_dpa4c_long_range*.pt2 \
+			${INFER_SCRIPT_PATH}/deeppot_dpa4c_long_range*.pte \
+			${INFER_SCRIPT_PATH}/deeppot_dpa4c_long_range*.expected
+		env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_long_range.py
+		exit 0
+	fi
+	# Remove stale generated model files so they can't be accidentally reused
+	# if gen scripts change format or the code version changes.
+	rm -f ${INFER_SCRIPT_PATH}/*.pt2 ${INFER_SCRIPT_PATH}/*.pte
 	# Run gen scripts in parallel for faster model generation.
 	# Wait on each PID separately so any failure is caught by set -e.
 	env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_sea.py &
@@ -113,6 +159,11 @@ else:
 	PID8=$!
 	wait $PID7
 	wait $PID8
+	# Copies of the multi-rank DPA3 archives of gen_dpa3.py and gen_spin.py
+	# whose nested with-comm artifact is corrupt; without them the load-failure
+	# cases of the pt_expt backends (test_with_comm_load_failure_ptexpt.cc)
+	# skip.
+	env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_corrupt_with_comm.py
 
 	# Native-spin DPA4 graph archives (baseline + model-level pair
 	# exclusion). Without this the whole native-spin graph C++ suite
@@ -134,6 +185,11 @@ else:
 	# exports NO with-comm artifact (single-rank only by construction).
 	env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin_zbl.py &
 	PID14=$!
+	# Minimal native-spin DPA4C for LAMMPS pair_style dpa4spin (no
+	# fparam/aparam/chg_spin). CPU-safe; used by the magnetic-force
+	# convention behavioral tests.
+	env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_spin_lmp.py &
+	PID14B=$!
 	# Native-spin DPA4C: the archives above all ship a with-comm artifact
 	# and are therefore served by DeepSpinPTExpt, which leaves
 	# NativeSpinPTExpt -- the class the dispatch selects for a native-spin
@@ -144,14 +200,29 @@ else:
 	if [ "${DP_VARIANT}" = "cuda" ]; then
 		env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_spin.py &
 		PID15=$!
+		# Compressed DPA4C, an energy and a native-spin model, on the compact
+		# canonical lower: the only fixtures that reach the device-resident
+		# entry points LAMMPS + Kokkos calls, with the replay of repeated
+		# evaluations and a runtime charge state.
+		env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_canonical.py &
+		PID16=$!
 	fi
+	# DPA4C with a latent-charge long-range module: the uncompressed archive
+	# takes the all-pairs topology of the graph lower; for the CUDA variant the
+	# script also writes the compressed archive, whose compact canonical lower
+	# reads the node positions.
+	env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_long_range.py &
+	PID17=$!
 	wait $PID11
 	wait $PID12
 	wait $PID13
 	wait $PID14
+	wait $PID14B
 	if [ "${DP_VARIANT}" = "cuda" ]; then
 		wait $PID15
+		wait $PID16
 	fi
+	wait $PID17
 fi
 if [ "${ENABLE_PADDLE:-TRUE}" == "TRUE" ]; then
 	PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir
diff --git a/source/install/test_lammps_kokkos.sh b/source/install/test_lammps_kokkos.sh
new file mode 100755
index 0000000000..26467343e2
--- /dev/null
+++ b/source/install/test_lammps_kokkos.sh
@@ -0,0 +1,116 @@
+#!/bin/bash
+set -euxo pipefail
+
+# Run after test_cc_local.sh (or --build-only then --models-only=kokkos):
+# reuse its C API/backend installation, long-range fixtures, and matching
+# LAMMPS checkout instead of rebuilding DeePMD or using the pip wheel.
+SCRIPT_PATH=$(dirname "$(realpath -s "$0")")
+REPO_ROOT=$(realpath "${SCRIPT_PATH}/../..")
+DEEPMD_C_ROOT=${DEEPMD_C_ROOT:-${REPO_ROOT}/dp_test}
+LAMMPS_SOURCE_ROOT=${LAMMPS_SOURCE_ROOT:-${REPO_ROOT}/source/build_tests/_deps/lammps_download-src}
+BUILD_ROOT=${REPO_ROOT}/source/build_lammps_kokkos
+LAMMPS_BUILD_DIR=${BUILD_ROOT}/lammps
+PLUGIN_BUILD_DIR=${BUILD_ROOT}/plugin
+PLUGIN_DIR=${BUILD_ROOT}/autoload
+NVCC_WRAPPER=${LAMMPS_SOURCE_ROOT}/lib/kokkos/bin/nvcc_wrapper
+# CUDA compilation can use substantial memory per process.
+BUILD_JOBS=${CMAKE_BUILD_PARALLEL_LEVEL:-4}
+
+if [[ ${DP_VARIANT:-} != cuda ]]; then
+	echo "This regression requires DP_VARIANT=cuda." >&2
+	exit 1
+fi
+test -f "${DEEPMD_C_ROOT}/lib/libdeepmd_c.so"
+test -f "${LAMMPS_SOURCE_ROOT}/cmake/CMakeLists.txt"
+test -x "${NVCC_WRAPPER}"
+command -v nvcc
+python -c 'import torch; assert torch.cuda.is_available(), "CUDA device is required"'
+
+# Use one MPI prefix for headers, libraries, Python bindings and launcher.
+# The CUDA job installs the MPICH wheel before building mpi4py and DeePMD.
+MPI_ROOT=${MPI_HOME:-$(python -c 'import sys; print(sys.prefix)')}
+MPI_C_COMPILER=$(command -v "${MPI_ROOT}/bin/mpicc")
+MPI_CXX_COMPILER=$(command -v "${MPI_ROOT}/bin/mpicxx")
+export MPI4PY_MPIABI=mpich
+export LD_LIBRARY_PATH=${MPI_ROOT}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+export DPA4SPIN_KOKKOS_MPIEXEC
+DPA4SPIN_KOKKOS_MPIEXEC=$(command -v "${MPI_ROOT}/bin/mpiexec")
+MPI_ARGS=(
+	"-DMPI_C_COMPILER=${MPI_C_COMPILER}"
+	"-DMPI_CXX_COMPILER=${MPI_CXX_COMPILER}"
+)
+
+# Kokkos detects the GPU architecture on the runner. Both builds use its
+# nvcc_wrapper and the plugin imports the exact Kokkos build configuration.
+cmake -S "${LAMMPS_SOURCE_ROOT}/cmake" -B "${LAMMPS_BUILD_DIR}" \
+	-D CMAKE_BUILD_TYPE=Release \
+	-D CMAKE_CXX_COMPILER="${NVCC_WRAPPER}" \
+	-D CMAKE_CXX_STANDARD=17 \
+	-D BUILD_SHARED_LIBS=ON \
+	-D BUILD_MPI=ON \
+	-D BUILD_OMP=OFF \
+	-D PKG_KOKKOS=ON \
+	-D Kokkos_ENABLE_CUDA=ON \
+	-D Kokkos_ENABLE_SERIAL=ON \
+	-D PKG_SPIN=ON \
+	-D PKG_PLUGIN=ON \
+	-D PKG_KSPACE=ON \
+	-D PKG_EXTRA-FIX=ON \
+	-D FFT=KISS \
+	-D FFT_KOKKOS=KISS \
+	"${MPI_ARGS[@]}"
+cmake --build "${LAMMPS_BUILD_DIR}" --target lammps --parallel "${BUILD_JOBS}"
+
+cmake -S "${REPO_ROOT}/source" -B "${PLUGIN_BUILD_DIR}" \
+	-D CMAKE_BUILD_TYPE=Release \
+	-D CMAKE_CXX_COMPILER="${NVCC_WRAPPER}" \
+	-D CMAKE_CXX_STANDARD=17 \
+	-D CMAKE_DISABLE_PRECOMPILE_HEADERS=ON \
+	-D BUILD_TESTING=OFF \
+	-D DEEPMD_C_ROOT="${DEEPMD_C_ROOT}" \
+	-D DP_USING_C_API=ON \
+	-D LAMMPS_SOURCE_ROOT="${LAMMPS_SOURCE_ROOT}" \
+	-D DEEPMD_LAMMPS_KOKKOS=ON \
+	-D Kokkos_DIR="${LAMMPS_BUILD_DIR}/lib/kokkos" \
+	"${MPI_ARGS[@]}"
+cmake --build "${PLUGIN_BUILD_DIR}" --target deepmd_lmp --parallel "${BUILD_JOBS}"
+mkdir -p "${PLUGIN_DIR}"
+ln -sfn "${PLUGIN_BUILD_DIR}/lmp/plugin/libdeepmd_lmp.so" "${PLUGIN_DIR}/dpplugin.so"
+
+export LD_LIBRARY_PATH=${LAMMPS_BUILD_DIR}:${DEEPMD_C_ROOT}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+python "${REPO_ROOT}/source/tests/infer/gen_dpa4c_spin_lmp_kokkos.py"
+
+# The wheel's Python package loads its own liblammps by absolute path. Merely
+# setting LD_LIBRARY_PATH would therefore still test the non-Kokkos wheel.
+export PYTHONPATH=${LAMMPS_SOURCE_ROOT}/python${PYTHONPATH:+:${PYTHONPATH}}
+export LAMMPS_PLUGIN_PATH=${PLUGIN_DIR}
+python - "${LAMMPS_SOURCE_ROOT}" "${LAMMPS_BUILD_DIR}" <<'PY'
+import sys
+from pathlib import Path
+
+import lammps
+
+source, build = map(Path, sys.argv[1:])
+assert Path(lammps.__file__).resolve() == (
+    source / "python/lammps/__init__.py"
+).resolve(), lammps.__file__
+with lammps.lammps(cmdargs=["-log", "none", "-screen", "none"]) as lmp:
+    mapped_libraries = {
+        Path(line.split()[-1]).resolve()
+        for line in Path("/proc/self/maps").read_text().splitlines()
+        if "/liblammps" in line
+    }
+    assert mapped_libraries == {(build / "liblammps.so").resolve()}, mapped_libraries
+    assert lmp.has_style("pair", "dpa4spin/kk"), "Kokkos plugin was not loaded"
+    print(f"Verified matching LAMMPS Python and shared library: {mapped_libraries}")
+PY
+
+DPA4SPIN_KOKKOS_REQUIRED=1 python -m pytest -v -s \
+	"${REPO_ROOT}/source/lmp/tests/test_dpa4spin_kokkos_magnetic_force.py"
+# The compressed long-range fixtures come from gen_dpa4c_long_range.py of
+# test_cc_local.sh; the MPI runner of the test launches mpirun from PATH, which
+# leads with the launcher of the MPI prefix, while pytest runs on the Python of
+# the job.
+PYTHON=$(command -v python)
+PATH=${MPI_ROOT}/bin:${PATH} DPA4C_LONG_RANGE_KOKKOS_REQUIRED=1 "${PYTHON}" -m pytest -v -s \
+	"${REPO_ROOT}/source/lmp/tests/test_lammps_dpa4c_long_range_kokkos.py"
diff --git a/source/lib/src/gpu/prod_force.cu b/source/lib/src/gpu/prod_force.cu
index 7b1359b3b0..993b6650cf 100644
--- a/source/lib/src/gpu/prod_force.cu
+++ b/source/lib/src/gpu/prod_force.cu
@@ -113,7 +113,14 @@ void prod_force_a_gpu(FPTYPE* force,
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
   const int ndescrpt = nnei * 4;
-  DPErrcheck(gpuMemset(force, 0, sizeof(FPTYPE) * nframes * nall * 3));
+  const size_t output_size =
+      sizeof(FPTYPE) * static_cast(nframes) * nall * 3;
+  if (output_size > 0) {
+    DPErrcheck(gpuMemset(force, 0, output_size));
+  }
+  if (nframes == 0 || nloc == 0 || nnei == 0) {
+    return;
+  }
 
   force_deriv_wrt_center_atom<<>>(
       force, net_deriv, in_deriv, ndescrpt, nloc, nall);
@@ -142,7 +149,14 @@ void prod_force_r_gpu(FPTYPE* force,
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
   const int ndescrpt = nnei * 1;
-  DPErrcheck(gpuMemset(force, 0, sizeof(FPTYPE) * nframes * nall * 3));
+  const size_t output_size =
+      sizeof(FPTYPE) * static_cast(nframes) * nall * 3;
+  if (output_size > 0) {
+    DPErrcheck(gpuMemset(force, 0, output_size));
+  }
+  if (nframes == 0 || nloc == 0 || nnei == 0) {
+    return;
+  }
 
   force_deriv_wrt_center_atom<<>>(
       force, net_deriv, in_deriv, ndescrpt, nloc, nall);
diff --git a/source/lib/src/gpu/prod_force_grad.cu b/source/lib/src/gpu/prod_force_grad.cu
index c784d6ba65..bcda4d2b7b 100644
--- a/source/lib/src/gpu/prod_force_grad.cu
+++ b/source/lib/src/gpu/prod_force_grad.cu
@@ -91,8 +91,12 @@ void prod_force_grad_a_gpu(FPTYPE* grad_net,
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
   const int ndescrpt = nnei * 4;
-  DPErrcheck(
-      gpuMemset(grad_net, 0, sizeof(FPTYPE) * nframes * nloc * ndescrpt));
+  const size_t output_size =
+      sizeof(FPTYPE) * static_cast(nframes) * nloc * ndescrpt;
+  if (output_size == 0) {
+    return;
+  }
+  DPErrcheck(gpuMemset(grad_net, 0, output_size));
   const int nblock = (ndescrpt + TPB - 1) / TPB;
   dim3 block_grid(nframes * nloc, nblock);
   dim3 thread_grid(TPB, 1);
@@ -122,8 +126,12 @@ void prod_force_grad_r_gpu(FPTYPE* grad_net,
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
   const int ndescrpt = nnei * 1;
-  DPErrcheck(
-      gpuMemset(grad_net, 0, sizeof(FPTYPE) * nframes * nloc * ndescrpt));
+  const size_t output_size =
+      sizeof(FPTYPE) * static_cast(nframes) * nloc * ndescrpt;
+  if (output_size == 0) {
+    return;
+  }
+  DPErrcheck(gpuMemset(grad_net, 0, output_size));
   const int nblock = (ndescrpt + TPB - 1) / TPB;
   dim3 block_grid(nframes * nloc, nblock);
   dim3 thread_grid(TPB, 1);
diff --git a/source/lib/src/gpu/prod_virial.cu b/source/lib/src/gpu/prod_virial.cu
index ab9c5326e3..55b9353922 100644
--- a/source/lib/src/gpu/prod_virial.cu
+++ b/source/lib/src/gpu/prod_virial.cu
@@ -116,7 +116,12 @@ void prod_virial_a_gpu(FPTYPE* virial,
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
   DPErrcheck(gpuMemset(virial, 0, sizeof(FPTYPE) * 9));
-  DPErrcheck(gpuMemset(atom_virial, 0, sizeof(FPTYPE) * 9 * nall));
+  if (nall > 0) {
+    DPErrcheck(gpuMemset(atom_virial, 0, sizeof(FPTYPE) * 9 * nall));
+  }
+  if (nloc == 0 || nnei == 0) {
+    return;
+  }
 
   const int LEN = 16;
   int nblock = (nnei + LEN - 1) / LEN;
@@ -146,7 +151,12 @@ void prod_virial_r_gpu(FPTYPE* virial,
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
   DPErrcheck(gpuMemset(virial, 0, sizeof(FPTYPE) * 9));
-  DPErrcheck(gpuMemset(atom_virial, 0, sizeof(FPTYPE) * 9 * nall));
+  if (nall > 0) {
+    DPErrcheck(gpuMemset(atom_virial, 0, sizeof(FPTYPE) * 9 * nall));
+  }
+  if (nloc == 0 || nnei == 0) {
+    return;
+  }
 
   const int LEN = 16;
   int nblock = (nnei + LEN - 1) / LEN;
diff --git a/source/lib/src/gpu/prod_virial_grad.cu b/source/lib/src/gpu/prod_virial_grad.cu
index dac5b20ba8..65471f0402 100644
--- a/source/lib/src/gpu/prod_virial_grad.cu
+++ b/source/lib/src/gpu/prod_virial_grad.cu
@@ -95,7 +95,12 @@ void prod_virial_grad_a_gpu(FPTYPE* grad_net,
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
   const int ndescrpt = nnei * 4;
-  DPErrcheck(gpuMemset(grad_net, 0, sizeof(FPTYPE) * nloc * ndescrpt));
+  const size_t output_size =
+      sizeof(FPTYPE) * static_cast(nloc) * ndescrpt;
+  if (output_size == 0) {
+    return;
+  }
+  DPErrcheck(gpuMemset(grad_net, 0, output_size));
   const int LEN = 128;
   const int nblock = (nloc + LEN - 1) / LEN;
   dim3 block_grid(nblock, nnei);
@@ -117,7 +122,12 @@ void prod_virial_grad_r_gpu(FPTYPE* grad_net,
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
   const int ndescrpt = nnei;
-  DPErrcheck(gpuMemset(grad_net, 0, sizeof(FPTYPE) * nloc * ndescrpt));
+  const size_t output_size =
+      sizeof(FPTYPE) * static_cast(nloc) * ndescrpt;
+  if (output_size == 0) {
+    return;
+  }
+  DPErrcheck(gpuMemset(grad_net, 0, output_size));
   const int LEN = 128;
   const int nblock = (nloc + LEN - 1) / LEN;
   dim3 block_grid(nblock, nnei);
diff --git a/source/lmp/compact_canonical_graph_kokkos.h b/source/lmp/compact_canonical_graph_kokkos.h
index 38389d6bd5..71321ebc04 100644
--- a/source/lmp/compact_canonical_graph_kokkos.h
+++ b/source/lmp/compact_canonical_graph_kokkos.h
@@ -6,15 +6,20 @@
 #ifndef LMP_COMPACT_CANONICAL_GRAPH_KOKKOS_H
 #define LMP_COMPACT_CANONICAL_GRAPH_KOKKOS_H
 
+#include 
 #include 
+#include 
 #include 
 #include 
 #include 
+#include 
+#include 
 #include 
 
 #include "atom.h"
 #include "atom_kokkos.h"
 #include "atom_masks.h"
+#include "comm_kokkos.h"
 #include "error.h"
 #include "kokkos_type.h"
 #include "neigh_list_kokkos.h"
@@ -23,15 +28,941 @@
 
 namespace LAMMPS_NS {
 
+#if defined(KOKKOS_ENABLE_CUDA)
+/// Arrangement of the 32 * K elements a warp holds K per lane in registers:
+/// element ``k * 32 + lane`` in register ``k`` of ``lane`` (register-major),
+/// or element ``lane * K + k`` (lane-major).
+enum class CompactCanonicalLayout { kRegisterMajor, kLaneMajor };
+
+// Bitonic sort of the 32 * K keys a warp holds K per lane, in the element
+// order of ``Layout``. Partners whose elements differ only in the bits of the
+// register index live in other registers of the same lane; the others live in
+// other lanes and are exchanged by shuffle. In the register-major layout the
+// five low bits of an element are those of its lane, so every stage of a
+// stride below the warp width shuffles; in the lane-major layout they are the
+// high bits, so only the stages of strides of at least ``K`` do, fifteen in
+// all, and the network runs without shared memory or barriers in either.
+// Every loop bound is a compile-time constant, which keeps the register array
+// statically indexed.
+template 
+__device__ __forceinline__ void compact_canonical_warp_sort(Key (&key)[K],
+                                                            const int lane) {
+  constexpr bool lane_major = Layout == CompactCanonicalLayout::kLaneMajor;
+  const auto element = [lane](const int k) {
+    return lane_major ? lane * K + k : k * 32 + lane;
+  };
+#pragma unroll
+  for (int size = 2; size <= 32 * K; size <<= 1) {
+#pragma unroll
+    for (int stride = size >> 1; stride > 0; stride >>= 1) {
+      if (lane_major ? stride < K : stride >= 32) {
+        const int register_stride = lane_major ? stride : stride >> 5;
+#pragma unroll
+        for (int k = 0; k < K; ++k) {
+          if ((k & register_stride) == 0) {
+            const int partner = k | register_stride;
+            const bool ascending = (element(k) & size) == 0;
+            const Key low = min(key[k], key[partner]);
+            const Key high = max(key[k], key[partner]);
+            key[k] = ascending ? low : high;
+            key[partner] = ascending ? high : low;
+          }
+        }
+      } else {
+        const int lane_stride = lane_major ? stride / K : stride;
+        const bool lower = (lane & lane_stride) == 0;
+#pragma unroll
+        for (int k = 0; k < K; ++k) {
+          const Key other = __shfl_xor_sync(0xffffffffu, key[k], lane_stride);
+          const bool ascending = (element(k) & size) == 0;
+          key[k] =
+              (lower == ascending) ? min(key[k], other) : max(key[k], other);
+        }
+      }
+    }
+  }
+}
+
+#endif
+
+/// Bits of the candidate slot in the sort keys of the single-pass fill, which
+/// serves candidate lists of up to ``1 << kCompactCanonicalSlotBits`` entries.
+constexpr int kCompactCanonicalSlotBits = 9;
+
+#if defined(KOKKOS_ENABLE_CUDA)
+// Read the packed record of one atom: its fixed-point coordinates and, in the
+// fourth slot, the model node it resolves to as a candidate. A record is 16
+// bytes and aligned to them, so it arrives in one 128-bit load.
+template 
+__device__ __forceinline__ int4 compact_canonical_record(const Packed& packed,
+                                                         const int atom) {
+  return __ldg(reinterpret_cast(&packed(atom, 0)));
+}
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+// Copy one word from global into shared memory without passing it through a
+// register, and fetch the whole 128-byte line that holds it into L2. A line
+// of the neighbor list holds one slot of 32 consecutive centers, which four
+// consecutive blocks of the fill stage at about the same time, so each line
+// is read from DRAM once and in one piece. The prefetch size takes PTX ISA
+// 7.4, CUDA 11.4; an earlier toolkit copies the word alone. The copies of a
+// thread complete at compact_canonical_await_copies.
+__device__ __forceinline__ void compact_canonical_copy_async(
+    int* destination, const int* source) {
+  const unsigned address =
+      static_cast(__cvta_generic_to_shared(destination));
+#if CUDART_VERSION >= 11040
+  asm volatile(
+      "cp.async.ca.shared.global.L2::128B [%0], [%1], 4;\n" ::"r"(address),
+      "l"(source)
+      : "memory");
+#else
+  asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" ::"r"(address),
+               "l"(source)
+               : "memory");
+#endif
+}
+
+__device__ __forceinline__ void compact_canonical_await_copies() {
+  asm volatile("cp.async.wait_all;\n" ::: "memory");
+}
+#endif
+
+// One component of the displacement between two fixed-point coordinates, in
+// fixed-point steps. A coordinate is held modulo 2^32, so the unsigned
+// difference wraps to the exact displacement whenever that displacement is
+// shorter than 2^31 steps, which the step of the fill guarantees for every
+// candidate.
+__device__ __forceinline__ float compact_canonical_displacement(const int qj,
+                                                                const int qi) {
+  return static_cast(
+      static_cast(static_cast(qj) - static_cast(qi)));
+}
+
+// Exclusive prefix of the published totals of every block before ``block``,
+// by the decoupled look-back of a single-pass prefix scan. The warp reads the
+// status words of 32 predecessors at once, nearest first, and sums the totals
+// down to the nearest inclusive prefix. Zero marks an unpublished block;
+// partial totals are stored plus one, and the high bit marks an inclusive
+// prefix whose remaining 63 bits cover the nonnegative int64 edge axis.
+__device__ __forceinline__ unsigned long long compact_canonical_look_back(
+    const int block, const int lane, unsigned long long* status) {
+  constexpr unsigned kFull = 0xffffffffu;
+  constexpr unsigned long long kInclusive = 1ull << 63;
+  unsigned long long prefix = 0;
+  for (int base = block - 1; base >= 0;) {
+    const int predecessor = base - lane;
+    unsigned long long value = kInclusive;
+    if (predecessor >= 0) {
+      value =
+          *(static_cast(status) + predecessor);
+    }
+    const unsigned inclusive = __ballot_sync(kFull, (value & kInclusive) != 0);
+    const unsigned pending = __ballot_sync(kFull, value == 0ull);
+    const int stop = inclusive ? __ffs(inclusive) - 1 : 31;
+    const unsigned needed = stop == 31 ? kFull : ((2u << stop) - 1u);
+    if (pending & needed) {
+      continue;
+    }
+    unsigned long long part =
+        lane <= stop
+            ? ((value & kInclusive) ? value & ~kInclusive : value - 1ull)
+            : 0ull;
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      part += __shfl_xor_sync(kFull, part, offset);
+    }
+    prefix += part;
+    if (inclusive) {
+      break;
+    }
+    base -= 32;
+  }
+  return prefix;
+}
+
+// The place of one block in the decoupled look-back of the single-pass build:
+// its index, the status words of every block, the total of the block, and the
+// shared words through which the warp that resolves the block offset claims
+// the look-back and hands the offset to the other warps. The total and the
+// inclusive prefix of a block may reach its status word in either order, since
+// different warps may store them: the total is exact and the inclusive prefix
+// is the sum of every total up to the block, so a later block that reads
+// either derives the same offset, and one that reads the total merely looks
+// further back. The pack kernel clears every status word before each fill
+// attempt on the same execution stream.
+struct CompactCanonicalScan {
+  int block;
+  unsigned long long* status;
+  unsigned total;
+  int* claim;
+  volatile unsigned long long* offset;
+  volatile int* ready;
+
+  // Look back for the block offset, publish the inclusive prefix of the block
+  // and hand the offset over, unless another warp of the block has claimed
+  // the look-back before; the first warp to arrive runs it.
+  __device__ __forceinline__ void resolve_once(const int lane) const {
+    constexpr unsigned kFull = 0xffffffffu;
+    constexpr unsigned long long kInclusive = 1ull << 63;
+    const int claimed = lane == 0 ? atomicExch(claim, 1) : 0;
+    if (__shfl_sync(kFull, claimed, 0) != 0) {
+      return;
+    }
+    const unsigned long long prefix =
+        compact_canonical_look_back(block, lane, status);
+    if (lane == 0) {
+      if (block != 0) {
+        const unsigned long long inclusive = prefix + total;
+        *(static_cast(status) + block) =
+            kInclusive | inclusive;
+      }
+      *offset = prefix;
+      __threadfence_block();
+      *ready = 1;
+    }
+  }
+
+  // The block offset, once ``resolve_once`` has handed it over. A device of
+  // compute capability below 7.0 has no sleep instruction and polls the
+  // volatile flag without pausing.
+  __device__ __forceinline__ std::int64_t await() const {
+    while (*ready == 0) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+      __nanosleep(32);
+#endif
+    }
+    __threadfence_block();
+    return static_cast(*offset);
+  }
+};
+
+// Order, place and emit one destination row of the single-pass build from the
+// keys of its candidates, ``K`` registers per lane, of which ``count`` are
+// survivors. The survivors are sorted on ``KS`` registers per lane, a power of
+// two: when that is fewer than ``K``, they are first compacted through the
+// warp's ``staging`` row of shared memory, and otherwise the keys stay in
+// their registers and the registers past ``K`` hold rejected keys, which sort
+// behind every survivor. Either way the multiset of keys, and therefore the
+// sorted row, is unchanged, while the sorting network shrinks with the square
+// of its depth. The network sorts in the lane-major layout, so that only its
+// fifteen stages across lanes shuffle, and the sorted row returns to the
+// register-major layout through ``slot_atom``, which holds the neighbor-list
+// entry of every candidate slot of the row until the source atoms of the
+// survivors take their place, at their positions in the row. The emitting
+// stores of a warp then cover consecutive positions.
+//
+// The row offset is the block offset plus ``row_offset`` within the block.
+// The first warp of the block to have sorted and laid out its row looks back
+// for the block offset, by which time the blocks before it have had the time
+// of the sort to publish their totals, and the other warps await it.
+template 
+__device__ __forceinline__ void compact_canonical_finish_row(
+    const std::uint32_t (&candidate_key)[K],
+    const unsigned count,
+    const unsigned row_offset,
+    const CompactCanonicalScan& scan,
+    const int lane,
+    std::uint32_t* staging,
+    int* slot_atom,
+    const int center,
+    const int mi,
+    const int4 position_i,
+    const int inum,
+    const int tail_begin,
+    const int tail_end,
+    const std::int64_t capacity,
+    std::int64_t* fill_total,
+    const float vector_scale,
+    const Tags& tag,
+    const Packed& packed,
+    const RowPointers& destination_row_ptr,
+    const Sources& source,
+    const Vectors& edge_vec,
+    const bool count_sources,
+    const Counts& source_counts) {
+  constexpr unsigned kFull = 0xffffffffu;
+  constexpr std::uint32_t kRejected = ~std::uint32_t{0};
+  constexpr int kSlot = kCompactCanonicalSlotBits;
+  constexpr std::uint32_t kSlotMask = (1u << kSlot) - 1u;
+  constexpr auto kLaneMajor = CompactCanonicalLayout::kLaneMajor;
+
+  // === Gather the survivors onto the sorting registers ===
+  // The network sorts any arrangement of its input, so the survivors enter it
+  // in the order they are gathered.
+  static_assert((KS & (KS - 1)) == 0 && KS >= 4,
+                "the sorting network spans 32 * KS keys, KS a power of two "
+                "of at least four");
+  std::uint32_t key[KS];
+  if constexpr (KS >= K) {
+#pragma unroll
+    for (int k = 0; k < KS; ++k) {
+      key[k] = k < K ? candidate_key[k] : kRejected;
+    }
+  } else {
+    const unsigned below = (1u << lane) - 1u;
+    unsigned placed = 0;
+#pragma unroll
+    for (int k = 0; k < K; ++k) {
+      const bool survivor = candidate_key[k] != kRejected;
+      const unsigned ballot = __ballot_sync(kFull, survivor);
+      if (survivor) {
+        staging[placed + __popc(ballot & below)] = candidate_key[k];
+      }
+      placed += __popc(ballot);
+    }
+    __syncwarp();
+#pragma unroll
+    for (int k = 0; k < KS; ++k) {
+      const unsigned position = k * 32 + lane;
+      key[k] = position < count ? staging[position] : kRejected;
+    }
+  }
+  compact_canonical_warp_sort(key, lane);
+
+  // === Order survivors of equal quantized length by source tag ===
+  // The predecessor of element ``lane * KS + k`` is the element before it in
+  // the same lane, or the last element of the previous lane.
+  bool tie = false;
+  const std::uint32_t last_above = __shfl_up_sync(kFull, key[KS - 1], 1);
+#pragma unroll
+  for (int k = 0; k < KS; ++k) {
+    const std::uint32_t previous =
+        k > 0 ? key[k > 0 ? k - 1 : 0] : (lane == 0 ? kRejected : last_above);
+    tie = tie || (key[k] != kRejected && previous != kRejected &&
+                  (key[k] >> kSlot) == (previous >> kSlot));
+  }
+  if (__any_sync(kFull, tie)) {
+    unsigned long long wide[KS];
+#pragma unroll
+    for (int k = 0; k < KS; ++k) {
+      wide[k] = ~0ull;
+      if (key[k] != kRejected) {
+        const int jj = static_cast(key[k] & kSlotMask);
+        const unsigned long long rank =
+            static_cast(tag(slot_atom[jj] & NEIGHMASK));
+        wide[k] =
+            (static_cast(key[k] >> kSlot) << (32 + kSlot)) |
+            (rank << kSlot) | static_cast(jj);
+      }
+    }
+    compact_canonical_warp_sort(wide, lane);
+#pragma unroll
+    for (int k = 0; k < KS; ++k) {
+      key[k] = static_cast(((wide[k] >> (32 + kSlot)) << kSlot) |
+                                          (wide[k] & kSlotMask));
+    }
+  }
+
+  // === Lay the source atoms out in the register-major order of the row ===
+  // The survivors precede the rejected keys, so position ``p`` of the row
+  // holds a survivor when ``p < count``. Every lane looks up the atoms of its
+  // positions before any is stored over the entries, and writes them four at
+  // a time: a quad that starts below ``count`` ends below ``32 * K``, inside
+  // the row.
+  int atom[KS];
+#pragma unroll
+  for (int k = 0; k < KS; ++k) {
+    atom[k] =
+        key[k] != kRejected ? slot_atom[key[k] & kSlotMask] & NEIGHMASK : 0;
+  }
+  __syncwarp();
+#pragma unroll
+  for (int quad = 0; quad < KS; quad += 4) {
+    if (lane * KS + quad < count) {
+      *reinterpret_cast(&slot_atom[lane * KS + quad]) =
+          make_int4(atom[quad], atom[quad + 1], atom[quad + 2], atom[quad + 3]);
+    }
+  }
+  __syncwarp();
+#pragma unroll
+  for (int k = 0; k < KS; ++k) {
+    const unsigned position = k * 32 + lane;
+    atom[k] = position < count ? slot_atom[position] : -1;
+  }
+
+  // === Await the block offset and publish the row ===
+  scan.resolve_once(lane);
+  const std::int64_t prefix = scan.await() + row_offset;
+  if (mi >= 0 && lane == 0) {
+    destination_row_ptr(mi) = prefix;
+  }
+  if (center == inum - 1) {
+    const std::int64_t total = prefix + count;
+    for (int m = tail_begin + lane; m <= tail_end; m += 32) {
+      destination_row_ptr(m) = total;
+    }
+    // Rows the previous build filled past this build's end become guards.
+    const std::int64_t end = total < capacity ? total : capacity;
+    const std::int64_t used = *fill_total;
+    for (std::int64_t position = end + lane; position < used; position += 32) {
+      source(position) = 0u;
+      edge_vec(3 * position + 0) = 0.0f;
+      edge_vec(3 * position + 1) = 0.0f;
+      edge_vec(3 * position + 2) = 0.0f;
+    }
+    __syncwarp();
+    if (lane == 0) {
+      *fill_total = end;
+    }
+  }
+
+  // === Emit the row in key order ===
+  // The records of four survivors are requested before the first of them is
+  // used, as in the key pass.
+  constexpr int kBatch = 4;
+  static_assert(KS % kBatch == 0, "the slots split into whole batches");
+#pragma unroll
+  for (int first = 0; first < KS; first += kBatch) {
+    int4 record[kBatch];
+#pragma unroll
+    for (int u = 0; u < kBatch; ++u) {
+      record[u] = atom[first + u] >= 0
+                      ? compact_canonical_record(packed, atom[first + u])
+                      : make_int4(0, 0, 0, 0);
+    }
+#pragma unroll
+    for (int u = 0; u < kBatch; ++u) {
+      const std::int64_t position = prefix + (first + u) * 32 + lane;
+      if (atom[first + u] >= 0 && position < capacity) {
+        source(position) = static_cast(record[u].w);
+        edge_vec(3 * position + 0) =
+            compact_canonical_displacement(record[u].x, position_i.x) *
+            vector_scale;
+        edge_vec(3 * position + 1) =
+            compact_canonical_displacement(record[u].y, position_i.y) *
+            vector_scale;
+        edge_vec(3 * position + 2) =
+            compact_canonical_displacement(record[u].z, position_i.z) *
+            vector_scale;
+        if (count_sources) {
+          Kokkos::atomic_fetch_add(&source_counts(record[u].w),
+                                   std::int64_t{1});
+        }
+      }
+    }
+  }
+}
+#endif
+
+// Fill the destination rows of ``Warps`` consecutive centers in one block of
+// the single-pass build, one warp per center.
+//
+// The rows are laid out in model-node order, which is the order of the
+// centers, so the row offset of a center is the number of edges of every
+// center before it. A block takes one ticket, which orders it after every
+// block that started before it, and fills the rows of the ``Warps`` centers
+// that the ticket names. Consecutive centers are spatial neighbors, so the
+// warps of a block gather largely the same neighbor coordinates and serve one
+// another from L1.
+//
+// The neighbor list stores the candidates of one slot for consecutive centers
+// next to each other, so the block reads the candidate entries of all its
+// centers together, eight centers of a slot per 32-byte sector, into shared
+// memory, where the warps read their own rows from then on. Each warp keys
+// its candidates; the block sums the survivors of its centers, publishes the
+// sum, and derives its offset by looking back over the published sums of the
+// blocks before it -- the decoupled look-back of a single-pass prefix scan --
+// once the first of its warps has sorted its row, while the others sort
+// theirs. A block publishes its sum as soon as it is known and its inclusive
+// prefix once its own look-back closes, so a window of 32 predecessors read
+// at once usually ends at an inclusive prefix. A block can wait only on
+// tickets taken before its own, which belong to running or finished blocks,
+// so the scan cannot deadlock.
+//
+// A status word is zero before publication, holds a block total plus one for
+// a partial sum, or holds an inclusive prefix with bit 63 set. The packing
+// kernel clears the status words and ticket counter before every fill attempt.
+// Edges beyond ``capacity`` are not written; the caller compares the total
+// with the capacity and repeats the build with larger arrays.
+//
+// The keys of up to ``32 * K`` candidates are formed in the registers of the
+// warp: the squared bond length quantized to the 23 bits above the nine bits
+// of the candidate slot, which resolves lengths far finer than any thermal
+// displacement. A rejected candidate carries the largest key and sorts behind
+// every survivor, so the sorted position of a survivor is its slot in the row;
+// the survivors are sorted on the fewest registers that hold them. The order
+// of the candidates in the neighbor list varies from build to build, and so
+// does the order of the local atoms, so survivors of equal quantized length
+// are ordered by the tag of their source atom: a row whose sorted keys show
+// such a tie is sorted again on 64-bit keys that insert the low 32 bits of the
+// tag between the length and the slot. Such ties arise in symmetric
+// configurations, such as a perfect lattice, and almost never in a thermal
+// one. A row holds every tag once when the box is thicker than twice the
+// cutoff along every periodic direction, as the folded node set of a single
+// rank requires, and the row order, and with it the summation order of every
+// reduction over the row, is then a function of the geometry and the atom
+// tags alone, which makes the forces reproducible from run to run. In a
+// thinner box, two images of one atom in a row tie on the tag as well and
+// keep the order of their slots.
+//
+// The warps also refresh the artifact's node types, which a center writes
+// for its own node and the halo nodes of the extended set take in a stride
+// over the centers, and the last center keeps the edge storage past the
+// physical end of the graph cleared: it zeroes the rows that the previous
+// build used and this one does not, up to the watermark ``fill_total``.
+//
+// Only a CUDA device pass compiles the body. The selection is made inside
+// this function rather than around its call because nvcc orders the captures
+// of an extended lambda by their first use in its body, so a body that
+// differs between the host and the device pass would name a different kernel
+// on each side.
+template 
+KOKKOS_INLINE_FUNCTION void compact_canonical_fill_row(
+    const int rank,
+    unsigned long long* ticket_counter,
+    unsigned long long* status,
+    const int inum,
+    const int tail_begin,
+    const int tail_end,
+    const std::int64_t capacity,
+    std::int64_t* fill_total,
+    const float cutsq,
+    const float key_scale,
+    const float vector_scale,
+    const Indices& numneigh,
+    const Neighbors& neighbors,
+    const ModelMap& loc2model,
+    const ModelMap& model2loc,
+    const Tags& tag,
+    const Packed& packed,
+    const Indices& type,
+    const TypeMap& type_map,
+    const RowPointers& model_type,
+    const RowPointers& destination_row_ptr,
+    const Sources& source,
+    const Vectors& edge_vec,
+    const bool count_sources,
+    const Counts& source_counts) {
+#if defined(KOKKOS_ENABLE_CUDA) && defined(__CUDA_ARCH__)
+  constexpr unsigned kFull = 0xffffffffu;
+  constexpr std::uint32_t kRejected = ~std::uint32_t{0};
+  // Highest length level of a survivor, one below the levels the key holds.
+  constexpr std::uint32_t kTopLevel =
+      (std::uint32_t{1} << (32 - kCompactCanonicalSlotBits)) - 2u;
+  constexpr unsigned long long kInclusive = 1ull << 63;
+  // Candidate slots of one row; the row stride adds four, which spreads the
+  // stores of a warp, eight centers by four slots, over all 32 banks.
+  constexpr int kSlots = 32 * K;
+  constexpr int kSlotStride = kSlots + 4;
+
+  // === Take the ticket that orders this block ===
+  const int lane = rank & 31;
+  const int warp = rank >> 5;
+  __shared__ unsigned long long block_ticket;
+  __shared__ unsigned center_count[Warps];
+  __shared__ unsigned long long block_offset;
+  __shared__ int claim;
+  __shared__ int ready;
+  __shared__ alignas(16) int slot_atom[Warps][kSlotStride];
+  __shared__ std::uint32_t staging[Warps][256];
+  if (rank == 0) {
+    block_ticket = atomicAdd(ticket_counter, 1ull);
+    claim = 0;
+    ready = 0;
+  }
+  __syncthreads();
+  const int block = static_cast(block_ticket);
+  const int center = block * Warps + warp;
+  const bool active = center < inum;
+  const int i = active ? center : 0;
+
+  // === Stage the candidate entries of every center of the block ===
+  // Thread ``rank`` stages slots ``rank / Warps + 32 s`` of center
+  // ``rank % Warps``, so a warp reads one sector per slot, holding the slot
+  // for all centers, and lane ``warp`` of every warp stages the warp's own
+  // center, whose model node and candidate count it hands the other lanes.
+  // The model nodes, the candidate counts and the center types are requested
+  // together. The entries are copied as the list holds them: every reader
+  // masks off the special-bond bits, and the key pass passes over the slots
+  // past the candidates of its center. Every copy of a thread is in flight
+  // before the first completes: on a device of compute capability 8.0 or
+  // higher as asynchronous copies, which hold no registers, and otherwise as
+  // loads that are all issued before the first is stored.
+  static_assert(32 % Warps == 0, "a warp covers whole slots of the block");
+  const int c = rank % Warps;
+  const int atom_c = block * Warps + c;
+  const int model_c = atom_c < inum ? loc2model(atom_c) : -1;
+  const int listed_c = atom_c < inum ? numneigh(atom_c) : 0;
+  const int type_i = active ? type(i) : 1;
+  const int candidates_c = model_c >= 0 ? listed_c : 0;
+  {
+    const int first_slot = rank / Warps;
+#if __CUDA_ARCH__ >= 800
+#pragma unroll
+    for (int s = 0; s < K; ++s) {
+      const int jj = first_slot + 32 * s;
+      if (jj < candidates_c) {
+        compact_canonical_copy_async(&slot_atom[c][jj], &neighbors(atom_c, jj));
+      }
+    }
+#else
+    int entry[K];
+#pragma unroll
+    for (int s = 0; s < K; ++s) {
+      const int jj = first_slot + 32 * s;
+      entry[s] = jj < candidates_c ? neighbors(atom_c, jj) : 0;
+    }
+#pragma unroll
+    for (int s = 0; s < K; ++s) {
+      const int jj = first_slot + 32 * s;
+      if (jj < candidates_c) {
+        slot_atom[c][jj] = entry[s];
+      }
+    }
+#endif
+  }
+  const int mi = __shfl_sync(kFull, model_c, warp);
+  const int candidates = __shfl_sync(kFull, candidates_c, warp);
+
+  // === Refresh the node types before the staged entries are awaited ===
+  if (mi >= 0 && lane == 0) {
+    model_type(mi) = type_map(type_i - 1);
+  }
+  if (active) {
+    for (std::int64_t m = tail_begin + std::int64_t{32} * center + lane;
+         m < tail_end; m += std::int64_t{32} * inum) {
+      model_type(m) = type_map(type(model2loc(m)) - 1);
+    }
+  }
+#if __CUDA_ARCH__ >= 800
+  compact_canonical_await_copies();
+#endif
+  __syncthreads();
+  const int* const row_atom = slot_atom[warp];
+  const int4 position_i = compact_canonical_record(packed, i);
+
+  // === Key every candidate and count the survivors ===
+  // The records of four slots are requested before the first of them is
+  // used, so the gathers of a lane overlap.
+  constexpr int kBatch = 4;
+  static_assert(K % kBatch == 0, "the slots split into whole batches");
+  std::uint32_t key[K];
+  unsigned count = 0;
+#pragma unroll
+  for (int first_slot = 0; first_slot < K; first_slot += kBatch) {
+    int4 record[kBatch];
+#pragma unroll
+    for (int u = 0; u < kBatch; ++u) {
+      const int slot = (first_slot + u) * 32 + lane;
+      record[u] = slot < candidates ? compact_canonical_record(
+                                          packed, row_atom[slot] & NEIGHMASK)
+                                    : make_int4(0, 0, 0, -1);
+    }
+#pragma unroll
+    for (int u = 0; u < kBatch; ++u) {
+      const int k = first_slot + u;
+      const float dx =
+          compact_canonical_displacement(record[u].x, position_i.x);
+      const float dy =
+          compact_canonical_displacement(record[u].y, position_i.y);
+      const float dz =
+          compact_canonical_displacement(record[u].z, position_i.z);
+      const float length2 = fmaf(dx, dx, fmaf(dy, dy, dz * dz));
+      // The level of a survivor stays below the top one, whatever the
+      // rounding of the scale, so that its key never equals the rejection
+      // value.
+      const std::uint32_t level =
+          min(static_cast(length2 * key_scale), kTopLevel);
+      key[k] = record[u].w >= 0 && length2 < cutsq
+                   ? (level << kCompactCanonicalSlotBits) |
+                         static_cast(k * 32 + lane)
+                   : kRejected;
+      count += __popc(__ballot_sync(kFull, key[k] != kRejected));
+    }
+  }
+  if (lane == 0) {
+    center_count[warp] = count;
+  }
+  __syncthreads();
+  unsigned row_offset = 0;
+  unsigned block_total = 0;
+#pragma unroll
+  for (int w = 0; w < Warps; ++w) {
+    const unsigned value = center_count[w];
+    row_offset += w < warp ? value : 0u;
+    block_total += value;
+  }
+
+  // === Publish the block total ===
+  if (warp == 0 && lane == 0) {
+    *(static_cast(status) + block) =
+        block == 0 ? kInclusive | block_total
+                   : static_cast(block_total) + 1ull;
+  }
+  if (!active) {
+    return;
+  }
+
+  // === Place, order and emit the row on the fewest registers that hold it ===
+  // Every block has the center of its first warp, so some warp of the block
+  // resolves the block offset.
+  const CompactCanonicalScan scan{block,  status,        block_total,
+                                  &claim, &block_offset, &ready};
+#define DEEPMD_COMPACT_CANONICAL_FINISH(KS)                               \
+  compact_canonical_finish_row(                                    \
+      key, count, row_offset, scan, lane, staging[warp], slot_atom[warp], \
+      center, mi, position_i, inum, tail_begin, tail_end, capacity,       \
+      fill_total, vector_scale, tag, packed, destination_row_ptr, source, \
+      edge_vec, count_sources, source_counts)
+  if (count <= 128u) {
+    DEEPMD_COMPACT_CANONICAL_FINISH(4);
+  } else if (K > 8 && count <= 256u) {
+    DEEPMD_COMPACT_CANONICAL_FINISH(8);
+  } else {
+    DEEPMD_COMPACT_CANONICAL_FINISH((K > 8 ? 16 : 8));
+  }
+#undef DEEPMD_COMPACT_CANONICAL_FINISH
+#endif
+}
+
+#if defined(KOKKOS_ENABLE_CUDA)
+// Sort the ``count`` distinct values of a row in place, ascending, with the
+// warp holding them ``K`` per lane in registers; the slots past the row hold
+// the largest value and sort behind it.
+template 
+__device__ __forceinline__ void compact_canonical_sort_in_registers(
+    std::int64_t* row, const int count, const int lane) {
+  std::int64_t key[K];
+#pragma unroll
+  for (int k = 0; k < K; ++k) {
+    const int index = k * 32 + lane;
+    key[k] = index < count ? row[index] : INT64_MAX;
+  }
+  compact_canonical_warp_sort(key,
+                                                                         lane);
+#pragma unroll
+  for (int k = 0; k < K; ++k) {
+    const int index = k * 32 + lane;
+    if (index < count) {
+      row[index] = key[k];
+    }
+  }
+}
+
+// Sort a row too long for the registers of a warp in place, ascending, by
+// odd-even transposition: ``count`` rounds of disjoint neighbour exchanges,
+// which the lanes of the warp share.
+__device__ inline void compact_canonical_sort_in_memory(std::int64_t* row,
+                                                        const long count,
+                                                        const int lane) {
+  for (long round = 0; round < count; ++round) {
+    for (long index = (round & 1) + 2 * lane; index + 1 < count; index += 64) {
+      const std::int64_t low = row[index];
+      const std::int64_t high = row[index + 1];
+      if (low > high) {
+        row[index] = high;
+        row[index + 1] = low;
+      }
+    }
+    __syncwarp();
+  }
+}
+#endif
+
+// Sort the rows of ``Warps`` consecutive nodes of the source CSR, one warp per
+// row, in the block of league rank ``league``. A row of up to 512 positions,
+// which holds every row of a node whose candidates the single-pass fill
+// serves, is sorted in the registers of its warp; a longer one in memory.
+//
+// Only a CUDA device pass compiles the body, for the reason given at
+// compact_canonical_fill_row.
+template 
+KOKKOS_INLINE_FUNCTION void compact_canonical_sort_source_rows(
+    const int rank,
+    const int league,
+    const int node_count,
+    const RowPointers& source_row_ptr,
+    const Order& source_order) {
+#if defined(KOKKOS_ENABLE_CUDA) && defined(__CUDA_ARCH__)
+  const int lane = rank & 31;
+  const int node = league * Warps + (rank >> 5);
+  if (node >= node_count) {
+    return;
+  }
+  const std::int64_t begin = source_row_ptr(node);
+  const std::int64_t count = source_row_ptr(node + 1) - begin;
+  std::int64_t* row = source_order.data() + begin;
+  if (count <= 1) {
+    return;
+  }
+  if (count <= 128) {
+    compact_canonical_sort_in_registers<4>(row, static_cast(count), lane);
+  } else if (count <= 256) {
+    compact_canonical_sort_in_registers<8>(row, static_cast(count), lane);
+  } else if (count <= 512) {
+    compact_canonical_sort_in_registers<16>(row, static_cast(count), lane);
+  } else {
+    compact_canonical_sort_in_memory(row, count, lane);
+  }
+#endif
+}
+
+// Accelerator stream of an execution space, which a device-resident model
+// evaluation is enqueued on so that it is ordered after the graph build and
+// before the force scatter without a host synchronization. Spaces without a
+// stream of their own yield null, which asks the model to synchronize.
+template 
+void* compact_canonical_stream(const ExecutionSpace&) {
+  return nullptr;
+}
+#if defined(KOKKOS_ENABLE_CUDA)
+inline void* compact_canonical_stream(const Kokkos::Cuda& space) {
+  return space.cuda_stream();
+}
+#endif
+
+// Pointers to the members of the swap plan of the Kokkos brick communication
+// -- the swap count, the atom counts, partners and first ghost slot of every
+// swap, and the device send lists -- which CommBrick and CommKokkos keep
+// protected. Formed through this derived class, they read the plan of any
+// CommKokkos; the class itself is never instantiated.
+struct CompactCanonicalSwapPlan : public CommKokkos {
+  static constexpr auto swap_count = &CompactCanonicalSwapPlan::nswap;
+  static constexpr auto send_count = &CompactCanonicalSwapPlan::sendnum;
+  static constexpr auto recv_count = &CompactCanonicalSwapPlan::recvnum;
+  static constexpr auto send_rank = &CompactCanonicalSwapPlan::sendproc;
+  static constexpr auto recv_rank = &CompactCanonicalSwapPlan::recvproc;
+  static constexpr auto first_recv = &CompactCanonicalSwapPlan::firstrecv;
+  static constexpr auto send_lists = &CompactCanonicalSwapPlan::k_sendlist;
+};
+
+// Fold of the per-atom outputs of the ghost nodes of an extended node set
+// onto their owners, ``size`` values per atom, in place of the host reverse
+// communication of the KOKKOS package. Wherever that package does not use its
+// device path for the reverse communication of a pair -- with MPI that cannot
+// address device memory, and under the classic or host settings of the pair
+// communication -- it sends the values through the host and copies the whole
+// per-atom arrays there and back at every call. The fold walks the swaps of
+// that communication in the same order instead: the pair packs the ghost
+// values of a swap on the device, only these values travel through pinned
+// host buffers to the rank that owns the atoms, and the pair adds them to its
+// atoms on the device. The atoms of the send list of a swap are distinct and
+// each receives one addition per swap, as in the host communication, so the
+// result is the same in every bit. Where the package uses its device path, or
+// with a communication style other than the Kokkos brick communication, it
+// communicates as it would.
+template 
+class CompactCanonicalGhostFold {
+  using Value = std::remove_const_t;
+
+ public:
+  template 
+  void fold(Comm* comm, MPI_Comm world, PairKokkos* pair, const int size) {
+    // The condition under which CommKokkos::reverse_comm(Pair*) takes its
+    // host path for a pair whose reverse communication runs on the device;
+    // LAMMPS names the flag "classic" before its 10Sep2025 release.
+    auto* kokkos_comm = dynamic_cast(comm);
+#if LAMMPS_VERSION_NUMBER < 20250910
+    const bool host_path =
+        kokkos_comm != nullptr && kokkos_comm->reverse_pair_comm_classic;
+#else
+    const bool host_path =
+        kokkos_comm != nullptr && kokkos_comm->reverse_pair_comm_legacy;
+#endif
+    if (!host_path) {
+      comm->reverse_comm(pair, size);
+      return;
+    }
+    using Plan = CompactCanonicalSwapPlan;
+    const int nswap = kokkos_comm->*Plan::swap_count;
+    const int* sendnum = kokkos_comm->*Plan::send_count;
+    const int* recvnum = kokkos_comm->*Plan::recv_count;
+    const int* sendproc = kokkos_comm->*Plan::send_rank;
+    const int* recvproc = kokkos_comm->*Plan::recv_rank;
+    const int* firstrecv = kokkos_comm->*Plan::first_recv;
+    auto& sendlists = kokkos_comm->*Plan::send_lists;
+    int capacity = 0;
+    for (int iswap = 0; iswap < nswap; ++iswap) {
+      capacity =
+          std::max(capacity, size * std::max(sendnum[iswap], recvnum[iswap]));
+    }
+    reserve(capacity);
+    sendlists.sync_device();
+    const MPI_Datatype type =
+        std::is_same::value ? MPI_FLOAT : MPI_DOUBLE;
+    auto d_send = send.template view();
+    auto d_recv = recv.template view();
+    for (int iswap = nswap - 1; iswap >= 0; --iswap) {
+      const int nsend = pair->pack_reverse_comm_kokkos(recvnum[iswap],
+                                                       firstrecv[iswap], send);
+      auto sendlist = Kokkos::subview(sendlists, iswap, Kokkos::ALL);
+      if (sendproc[iswap] == comm->me) {
+        pair->unpack_reverse_comm_kokkos(sendnum[iswap], sendlist, send);
+        continue;
+      }
+      const int nrecv = size * sendnum[iswap];
+      // A copy between views without an execution space waits for the device
+      // before and after it, so the packed values are complete when sent and
+      // the received ones are on the device before they are added.
+      Kokkos::deep_copy(Kokkos::subview(h_send, std::make_pair(0, nsend)),
+                        Kokkos::subview(d_send, std::make_pair(0, nsend)));
+      MPI_Request request;
+      if (sendnum[iswap] > 0) {
+        MPI_Irecv(h_recv.data(), nrecv, type, sendproc[iswap], 0, world,
+                  &request);
+      }
+      if (recvnum[iswap] > 0) {
+        MPI_Send(h_send.data(), nsend, type, recvproc[iswap], 0, world);
+      }
+      if (sendnum[iswap] > 0) {
+        MPI_Wait(&request, MPI_STATUS_IGNORE);
+      }
+      Kokkos::deep_copy(Kokkos::subview(d_recv, std::make_pair(0, nrecv)),
+                        Kokkos::subview(h_recv, std::make_pair(0, nrecv)));
+      pair->unpack_reverse_comm_kokkos(sendnum[iswap], sendlist, recv);
+    }
+  }
+
+ private:
+  void reserve(const int capacity) {
+    if (static_cast(send.extent(0)) >= capacity) {
+      return;
+    }
+    send = Buffer("compact_canonical:fold_send", capacity);
+    recv = Buffer("compact_canonical:fold_recv", capacity);
+    h_send = Kokkos::View(
+        "compact_canonical:fold_send_host", capacity);
+    h_recv = Kokkos::View(
+        "compact_canonical:fold_recv_host", capacity);
+  }
+
+  // The pair packs into and unpacks from the device views of the buffers;
+  // only the packed values of a swap are staged in the pinned host views.
+  Buffer send, recv;
+  Kokkos::View h_send, h_recv;
+};
+
 // Device-resident model node set and compact canonical neighbor graph.
 //
 // A ``.pt2`` artifact frozen with the compact canonical lower consumes a
-// dual-CSR topology over a contiguous range of model nodes: uint32 source
-// indices, float32 edge vectors, an int64 row pointer per node in each
-// direction, and a permutation that reorders the destination-major edge list
-// into source-major order. This component derives all of it from the Kokkos
-// device neighbor list and leaves it on the device, so a pair style hands the
-// model raw device pointers and never stages the graph through the host.
+// destination-major topology over a contiguous range of model nodes, with
+// 32-bit source-node indices, float32 edge vectors and 64-bit CSR positions.
+// The source-major CSR -- a second row pointer and the permutation that
+// reorders the edge list into source order -- is built only on request, for an
+// artifact whose force assembly groups edges by source or for an evaluation
+// that wants the source-attributed per-atom virial. This component derives
+// the graph from the Kokkos device neighbor list and leaves it on the device,
+// so a pair style hands the model raw device pointers and never stages the
+// graph through the host.
 //
 // The node set is the model's view of the atoms. Atom types that map to no
 // model element (NULL in ``pair_coeff``) are virtual: they carry no node, so
@@ -43,9 +974,7 @@ namespace LAMMPS_NS {
 // ghost a node of its own -- the extended node set -- and leaves the pair
 // style to fold the ghost outputs onto their owners by reverse communication.
 //
-// build() is one function, and public, because CUDA forbids extended device
-// lambdas inside non-public members: its stages cannot be split into private
-// helpers.
+// Functions that enclose extended device lambdas are public, as CUDA requires.
 template 
 class CompactCanonicalGraphKokkos : protected Pointers {
  public:
@@ -80,6 +1009,15 @@ class CompactCanonicalGraphKokkos : protected Pointers {
       h_type_map(t) = host_type_map[t];
     }
     Kokkos::deep_copy(d_type_map, h_type_map);
+    // The pair styles construct the builder before LAMMPS confirms that
+    // KOKKOS is enabled, so its views are allocated here, once; the fill
+    // watermark starts at zero, as fresh storage is.
+    if (!d_fill_total.is_allocated()) {
+      d_fill_total = Kokkos::View(
+          "compact_canonical:fill_total");
+      h_edge_count = Kokkos::View(
+          "compact_canonical:edge_count");
+    }
   }
 
   // Rebuild the atom -> node compaction, which only moves when the neighbor
@@ -157,16 +1095,39 @@ class CompactCanonicalGraphKokkos : protected Pointers {
   }
 
   // Build the compact canonical graph of the current configuration from the
-  // Kokkos full neighbor list. Bond vectors are the center-to-neighbor
-  // displacements divided by ``dist_unit_cvt_factor``, the same conversion the
-  // pair style applies to coordinates it hands the model.
+  // Kokkos full neighbor list, with the source-major CSR when ``source_csr``
+  // asks for it. Bond vectors are the center-to-neighbor displacements divided
+  // by ``dist_unit_cvt_factor``, the same conversion the pair style applies to
+  // coordinates it hands the model.
   void build(class NeighList* list,
              double cutoff,
-             double dist_unit_cvt_factor) {
+             double dist_unit_cvt_factor,
+             bool source_csr) {
     refresh_nodes();
-
     auto* k_list = static_cast*>(list);
     const int inum = k_list->inum;
+    // The candidate lists move only with the neighbor list: their longest
+    // entry selects the fill, and the candidates of the model centers bound
+    // the edge count until the next rebuild.
+    if (neighbor->ago == 0 || longest_candidates < 0) {
+      auto numneigh = k_list->d_numneigh;
+      auto ilist = k_list->d_ilist;
+      auto loc2model = d_loc2model;
+      int longest = 0;
+      std::int64_t total = 0;
+      Kokkos::parallel_reduce(
+          "compact_canonical:candidates",
+          Kokkos::RangePolicy(0, inum),
+          KOKKOS_LAMBDA(const int ii, int& value, std::int64_t& sum) {
+            const int i = ilist(ii);
+            const int count = numneigh(i);
+            value = count > value ? count : value;
+            sum += loc2model(i) >= 0 ? count : 0;
+          },
+          Kokkos::Max(longest), Kokkos::Sum(total));
+      longest_candidates = longest;
+      candidate_total = total;
+    }
     auto d_numneigh = k_list->d_numneigh;
     auto d_neighbors = k_list->d_neighbors;
     auto d_ilist = k_list->d_ilist;
@@ -180,6 +1141,8 @@ class CompactCanonicalGraphKokkos : protected Pointers {
     const int nall = atom->nlocal + atom->nghost;
 
     // === Node types in the artifact's index layout ===
+    // The single-pass fill writes them alongside the rows; every other
+    // configuration maps them here.
     atomKK->sync(execution_space, TYPE_MASK);
     auto type = atomKK->k_type.template view();
     auto type_map = d_type_map;
@@ -188,28 +1151,27 @@ class CompactCanonicalGraphKokkos : protected Pointers {
           "compact_canonical:model_type", nall);
     }
     auto model_type = d_model_type;
-    Kokkos::parallel_for(
-        "compact_canonical:node_type",
-        Kokkos::RangePolicy(0, node_count_int),
-        KOKKOS_LAMBDA(const int m) {
-          model_type(m) = type_map(type(model2loc(m)) - 1);
-        });
 
-    atomKK->sync(execution_space, X_MASK);
+    atomKK->sync(execution_space, X_MASK | TAG_MASK);
     auto x = atomKK->k_x.template view();
+    auto tag = atomKK->k_tag.template view();
 
     if (d_destination_row_ptr.extent(0) < node_count + 1) {
       d_destination_row_ptr = Kokkos::View(
           "compact_canonical:destination_row_ptr", node_count + 1);
-      d_source_counts = Kokkos::View(
+    }
+    if (source_csr && d_source_row_ptr.extent(0) < node_count + 1) {
+      d_source_counts = Kokkos::View(
           "compact_canonical:source_counts", node_count);
       d_source_row_ptr = Kokkos::View(
           "compact_canonical:source_row_ptr", node_count + 1);
-      d_source_cursor = Kokkos::View(
+      d_source_cursor = Kokkos::View(
           "compact_canonical:source_cursor", node_count);
     }
-    Kokkos::deep_copy(d_destination_row_ptr, std::int64_t{0});
-    Kokkos::deep_copy(d_source_counts, std::uint32_t{0});
+    if (source_csr) {
+      Kokkos::deep_copy(typename DeviceType::execution_space(), d_source_counts,
+                        std::int64_t{0});
+    }
     if (node_count_int == 0) {
       storage_count = 0;
       return;
@@ -217,165 +1179,360 @@ class CompactCanonicalGraphKokkos : protected Pointers {
     auto destination_row_ptr = d_destination_row_ptr;
     auto source_counts = d_source_counts;
 
-    // === Destination-major CSR: per-node edge count, then its prefix sum ===
-    Kokkos::parallel_for(
-        "compact_canonical:count", Kokkos::RangePolicy(0, inum),
-        KOKKOS_LAMBDA(const int ii) {
-          const int i = d_ilist(ii);
-          const int mi = loc2model(i);
-          if (mi < 0) {
-            return;
-          }
-          const double xi = x(i, 0);
-          const double yi = x(i, 1);
-          const double zi = x(i, 2);
-          const int jnum = d_numneigh(i);
-          std::int64_t count = 0;
-          for (int jj = 0; jj < jnum; ++jj) {
-            const int j = d_neighbors(i, jj) & NEIGHMASK;
-            const int mj = candidate_to_model(j);
-            if (mj < 0) {
-              continue;
-            }
-            const double dx = x(j, 0) - xi;
-            const double dy = x(j, 1) - yi;
-            const double dz = x(j, 2) - zi;
-            if (dx * dx + dy * dy + dz * dz < cutsq) {
-              ++count;
-            }
-          }
-          destination_row_ptr(mi) = count;
-        });
-
-    Kokkos::parallel_scan(
-        "compact_canonical:destination_scan",
-        Kokkos::RangePolicy(0, node_count_int),
-        KOKKOS_LAMBDA(const int node, std::int64_t& update, const bool final) {
-          const std::int64_t count = destination_row_ptr(node);
-          if (final) {
-            destination_row_ptr(node) = update;
-          }
-          update += count;
-          if (final && node == node_count_int - 1) {
-            destination_row_ptr(node_count_int) = update;
-          }
-        });
-    std::int64_t edge_count = 0;
-    Kokkos::deep_copy(edge_count,
-                      Kokkos::subview(d_destination_row_ptr, node_count_int));
-    storage_count = std::max(edge_count, min_storage_edges);
-    if (static_cast(storage_count) >
-        std::numeric_limits::max()) {
+    // === Destination-major fill, one warp per center, by bond length ===
+    // Every row is emitted in increasing bond length. A compressed descriptor
+    // reads one radial-table interval per edge, and rows of one length order
+    // make the concurrently resident warps of a multiprocessor, each scanning
+    // its own center, address the same intervals at the same time, so the
+    // table is served from L1 instead of being streamed from L2 once per edge.
+    // In a box thicker than twice the cutoff, where a row holds every tag
+    // once, the order follows from the geometry and the atom tags alone, which
+    // also fixes the summation order of every destination reduction
+    // independently of the neighbor list (see compact_canonical_fill_row).
+    //
+    // The team fill forms one 64-bit sort key per survivor: the squared length
+    // quantized to 16 bits, then the tag of the source atom, then the
+    // candidate slot. The keys are sorted and the survivors written at
+    // consecutive slots, so the dominant store stream stays coalesced; the
+    // geometry of a survivor is recomputed from its slot, whose coordinates
+    // the key pass has just brought into cache. The single-pass fill keys a
+    // survivor by its squared length quantized to 23 bits and its slot, and
+    // orders survivors of equal quantized length by tag only in the rows that
+    // hold such ties. The two orders differ.
+    //
+    // On a CUDA device whose candidate lists fit 512 entries, the warp of a
+    // center sorts in registers and the row offsets come from the decoupled
+    // look-back of compact_canonical_fill_row, so counting, offsets and the
+    // fill are one pass. Every other configuration counts the rows first,
+    // scans the counts, and fills from team scratch.
+    constexpr int neighbor_lanes = 32;  // lanes cooperating on one center
+    // Consecutive centers filled by one block of the single-pass fill: eight,
+    // so that one 32-byte sector of the neighbor list holds a candidate slot
+    // of every center of the block.
+    constexpr int fill_warps = 8;
+    static_assert(16 * neighbor_lanes <= (1 << kCompactCanonicalSlotBits),
+                  "the single-pass sort keys must address every candidate");
+    const int key_capacity = static_cast(d_neighbors.extent(1));
+    if (key_capacity >= (1 << 16)) {
       error->one(FLERR,
-                 "Compact canonical graph exceeds the uint32 edge-index range");
-    }
-    const std::size_t required = static_cast(storage_count);
-    if (edge_capacity < required) {
-      // Thermal cutoff-count fluctuations are much smaller than the historical
-      // 12.5% geometric-growth reserve. A 2% reserve avoids repeated allocation
-      // while preventing unused edge storage from retaining several GiB at
-      // billion-edge scale.
-      const std::size_t slack = required / 50 + 64;
-      if (required > std::numeric_limits::max() - slack) {
-        error->one(FLERR, "Compact canonical graph capacity overflows size_t");
-      }
-      edge_capacity = required + slack;
-      d_source = Kokkos::View(
-          "compact_canonical:source", edge_capacity);
-      d_edge_vec = Kokkos::View(
-          "compact_canonical:edge_vec", edge_capacity * 3);
-      d_source_order = Kokkos::View(
-          "compact_canonical:source_order", edge_capacity);
+                 "Compact canonical graph supports at most 65535 neighbor "
+                 "candidates per atom");
     }
-
-    auto source = d_source;
-    auto edge_vec = d_edge_vec;
-    // === Destination-major fill, one warp per center ===
-    // A thread-per-center fill writes each surviving edge at an offset private
-    // to its center, so the lanes of a warp scatter their twelve-byte edge
-    // vectors across thirty-two unrelated rows. Cooperating on one center
-    // instead sends consecutive survivors to consecutive slots, which coalesces
-    // the dominant store stream. Candidates are taken a warp at a time and each
-    // lane writes at its exclusive prefix within the warp, so the edge order is
-    // the candidate order a serial fill produces.
-    constexpr int neighbor_lanes = 32;  // lanes cooperating on one center
     using team_policy = Kokkos::TeamPolicy;
+    // The single-pass fill runs one warp per center and eight consecutive
+    // centers per team. A team that keys up to twelve candidates per lane
+    // stages them in about 21 KB of shared memory, so four teams, 32 warps,
+    // stay resident per multiprocessor and hide the latency of the staging
+    // and the gathers; one that keys sixteen takes 25 KB, and three teams, 24
+    // warps, fit. Either bound holds the registers of the whole kernel at the
+    // budget of its common path rather than at that of the rarely taken
+    // tie-break sort.
+    using fill_policy = Kokkos::TeamPolicy<
+        DeviceType,
+        Kokkos::LaunchBounds>;
+    using wide_fill_policy = Kokkos::TeamPolicy<
+        DeviceType,
+        Kokkos::LaunchBounds>;
     using member_type = typename team_policy::member_type;
-    using lane_scratch =
-        Kokkos::View>;
-    using vector_scratch =
-        Kokkos::View>;
-    const int scratch_bytes = lane_scratch::shmem_size(neighbor_lanes) +
-                              vector_scratch::shmem_size(3 * neighbor_lanes);
-    Kokkos::parallel_for(
-        "compact_canonical:fill",
-        team_policy(inum, neighbor_lanes)
-            .set_scratch_size(0, Kokkos::PerTeam(scratch_bytes)),
-        KOKKOS_LAMBDA(const member_type& team) {
-          const int i = d_ilist(team.league_rank());
-          const int mi = loc2model(i);
-          if (mi < 0) {
-            return;
-          }
-          lane_scratch node(team.team_scratch(0), neighbor_lanes);
-          vector_scratch vec(team.team_scratch(0), 3 * neighbor_lanes);
-
-          const double xi = x(i, 0);
-          const double yi = x(i, 1);
-          const double zi = x(i, 2);
-          const int jnum = d_numneigh(i);
-          const int lane = team.team_rank();
-          std::int64_t edge = destination_row_ptr(mi);
-          for (int base = 0; base < jnum; base += neighbor_lanes) {
-            const int jj = base + lane;
-            int mj = -1;
-            if (jj < jnum) {
+    // Squared lengths quantize to the key bits above the slot: 23 bits in the
+    // single-pass fill and 16 bits in the 64-bit keys of the team fill.
+    const double fill_key_levels =
+        static_cast((1u << (32 - kCompactCanonicalSlotBits)) - 1u);
+    const double key_scale = 65535.0 / cutsq;
+    std::int64_t edge_count = 0;
+    // The single-pass fill lays the rows out in the order of the centers and
+    // writes the row pointer and the type of a local node from its own center,
+    // so it serves a list whose centers are all the local atoms in index
+    // order, `ilist(ii) == ii`, as the Kokkos neighbor lists give them, and
+    // takes center `ii` for local atom `ii`; a list that leaves some out
+    // (`neigh_modify include`) takes the counted fill.
+    bool single_pass = false;
+#if defined(KOKKOS_ENABLE_CUDA)
+    single_pass =
+        std::is_same_v &&
+        inum > 0 && inum == atom->nlocal &&
+        longest_candidates <= 16 * neighbor_lanes;
+#endif
+    if (single_pass) {
+      // === Pack every atom the fill can visit into one aligned record ===
+      // The fill works on fixed-point coordinates, the low 32 bits of
+      // round(x / step) for a step of 2^-position_bits length units, the
+      // finest power of two for which 2^31 steps span four neighbor cutoffs.
+      // The wrapped difference of two coordinates is then the exact
+      // fixed-point displacement of any pair closer than four neighbor
+      // cutoffs, twice the farthest a candidate drifts from its center while
+      // its neighbor list is valid, so every bond vector and squared length of
+      // the fill is formed in single precision from integers. A bond vector
+      // carries, on top of its own single-precision rounding, the rounding of
+      // its two coordinates to the step: at most one step, 1.5e-8 length
+      // units at a neighbor cutoff of 8.
+      const int position_bits =
+          std::ilogb(std::ldexp(1.0, 29) / neighbor->cutneighmax);
+      const double position_scale = std::ldexp(1.0, position_bits);
+      const double step = std::ldexp(1.0, -position_bits);
+      const float vector_scale = static_cast(step * inv_dist);
+      if ((int)d_packed.extent(0) < nall) {
+        d_packed = Kokkos::View(
+            "compact_canonical:packed", nall);
+      }
+      auto packed = d_packed;
+      // The cutoff and the key scale in squared fixed-point steps. The scale
+      // derives from the rounded cutoff, which keeps every key of a squared
+      // length below that cutoff within the key bits.
+      const float fill_cutsq =
+          static_cast(cutsq * position_scale * position_scale);
+      const float fill_key_scale =
+          static_cast(fill_key_levels / fill_cutsq);
+      const int teams = (inum + fill_warps - 1) / fill_warps;
+      if ((int)d_fill_status.extent(0) < teams) {
+        d_fill_status = Kokkos::View(
+            "compact_canonical:fill_status", teams);
+      }
+      if (d_fill_ticket.extent(0) == 0) {
+        d_fill_ticket = Kokkos::View(
+            "compact_canonical:fill_ticket", 1);
+      }
+      auto status = d_fill_status;
+      auto ticket = d_fill_ticket;
+      // Edges are a subset of the candidates, so storage sized by the
+      // candidate total cannot overflow before the next rebuild and the edge
+      // count can stay on the device. Systems whose candidate storage exceeds
+      // the bound read the count back instead, which their step amortizes.
+      const bool bounded =
+          !source_csr && static_cast(candidate_total) <=
+                             bounded_storage_bytes() / edge_storage_bytes;
+      reserve_edges(bounded ? candidate_total
+                            : static_cast(edge_capacity));
+      for (;;) {
+        // Packing initializes the scan before each fill attempt, including a
+        // retry after growing edge storage. Every block owns at least one
+        // center, so the atom axis covers every active status word.
+        Kokkos::parallel_for(
+            "compact_canonical:pack", Kokkos::RangePolicy(0, nall),
+            KOKKOS_LAMBDA(const int a) {
+              for (int d = 0; d < 3; ++d) {
+                packed(a, d) = static_cast(
+                    static_cast(static_cast(
+                        Kokkos::nearbyint(x(a, d) * position_scale))));
+              }
+              packed(a, 3) = candidate_to_model(a);
+              if (a < teams) {
+                status(a) = 0ull;
+              }
+              if (a == 0) {
+                ticket(0) = 0ull;
+              }
+            });
+        const std::int64_t capacity = static_cast(edge_capacity);
+        if (longest_candidates <= 8 * neighbor_lanes) {
+          fill_rows<8, fill_warps>(
+              fill_policy(teams, fill_warps * neighbor_lanes), k_list,
+              fill_cutsq, fill_key_scale, vector_scale, source_csr);
+        } else if (longest_candidates <= 12 * neighbor_lanes) {
+          fill_rows<12, fill_warps>(
+              fill_policy(teams, fill_warps * neighbor_lanes), k_list,
+              fill_cutsq, fill_key_scale, vector_scale, source_csr);
+        } else {
+          fill_rows<16, fill_warps>(
+              wide_fill_policy(teams, fill_warps * neighbor_lanes), k_list,
+              fill_cutsq, fill_key_scale, vector_scale, source_csr);
+        }
+        if (bounded) {
+          break;
+        }
+        edge_count = read_edge_count(node_count_int);
+        if (edge_count <= capacity) {
+          break;
+        }
+        // Edges past the capacity were dropped and the source histogram is
+        // incomplete; the build repeats with arrays of the measured size.
+        reserve_edges(edge_count);
+        if (source_csr) {
+          Kokkos::deep_copy(typename DeviceType::execution_space(),
+                            d_source_counts, std::int64_t{0});
+        }
+      }
+    } else {
+      Kokkos::parallel_for(
+          "compact_canonical:node_type",
+          Kokkos::RangePolicy(0, node_count_int),
+          KOKKOS_LAMBDA(const int m) {
+            model_type(m) = type_map(type(model2loc(m)) - 1);
+          });
+
+      // === Per-node edge count, then its prefix sum ===
+      Kokkos::deep_copy(typename DeviceType::execution_space(),
+                        d_destination_row_ptr, std::int64_t{0});
+      Kokkos::parallel_for(
+          "compact_canonical:count", Kokkos::RangePolicy(0, inum),
+          KOKKOS_LAMBDA(const int ii) {
+            const int i = d_ilist(ii);
+            const int mi = loc2model(i);
+            if (mi < 0) {
+              return;
+            }
+            const double xi = x(i, 0);
+            const double yi = x(i, 1);
+            const double zi = x(i, 2);
+            const int jnum = d_numneigh(i);
+            std::int64_t count = 0;
+            for (int jj = 0; jj < jnum; ++jj) {
               const int j = d_neighbors(i, jj) & NEIGHMASK;
-              mj = candidate_to_model(j);
-              if (mj >= 0) {
-                const double dx = x(j, 0) - xi;
-                const double dy = x(j, 1) - yi;
-                const double dz = x(j, 2) - zi;
-                if (dx * dx + dy * dy + dz * dz < cutsq) {
-                  vec(3 * lane + 0) = static_cast(dx * inv_dist);
-                  vec(3 * lane + 1) = static_cast(dy * inv_dist);
-                  vec(3 * lane + 2) = static_cast(dz * inv_dist);
-                } else {
-                  mj = -1;
-                }
+              const int mj = candidate_to_model(j);
+              if (mj < 0) {
+                continue;
               }
+              const double dx = x(j, 0) - xi;
+              const double dy = x(j, 1) - yi;
+              const double dz = x(j, 2) - zi;
+              if (dx * dx + dy * dy + dz * dz < cutsq) {
+                ++count;
+              }
+            }
+            destination_row_ptr(mi) = count;
+          });
+      Kokkos::parallel_scan(
+          "compact_canonical:destination_scan",
+          Kokkos::RangePolicy(0, node_count_int),
+          KOKKOS_LAMBDA(const int node, std::int64_t& update,
+                        const bool final) {
+            const std::int64_t count = destination_row_ptr(node);
+            if (final) {
+              destination_row_ptr(node) = update;
             }
-            node(lane) = mj;
-            team.team_barrier();
-
-            // Compact the survivors of this warp of candidates: consecutive
-            // survivors take consecutive slots, so the stores of a warp fall in
-            // one contiguous span of the edge arrays.
-            std::int64_t kept = 0;
-            Kokkos::parallel_scan(
-                Kokkos::TeamThreadRange(team, neighbor_lanes),
-                [&](const int slot, std::int64_t& offset, const bool final) {
-                  const int target = node(slot);
-                  if (final && target >= 0) {
-                    const std::int64_t position = edge + offset;
-                    source(position) = static_cast(target);
-                    edge_vec(3 * position + 0) = vec(3 * slot + 0);
-                    edge_vec(3 * position + 1) = vec(3 * slot + 1);
-                    edge_vec(3 * position + 2) = vec(3 * slot + 2);
-                    Kokkos::atomic_fetch_add(&source_counts(target),
-                                             std::uint32_t{1});
+            update += count;
+            if (final && node == node_count_int - 1) {
+              destination_row_ptr(node_count_int) = update;
+            }
+          });
+      edge_count = read_edge_count(node_count_int);
+      reserve_edges(edge_count);
+
+      // === Fill from team scratch ===
+      auto source = d_source;
+      auto edge_vec = d_edge_vec;
+      using key_scratch =
+          Kokkos::View>;
+      const int scratch_bytes = key_scratch::shmem_size(neighbor_lanes) +
+                                key_scratch::shmem_size(key_capacity);
+      // The keys of a row stay in the on-chip scratch of level 0 while the
+      // candidate lists fit it, and in the device memory of level 1 beyond.
+      team_policy fill_teams(inum, neighbor_lanes);
+      const int scratch_level =
+          scratch_bytes <= fill_teams.scratch_size_max(0) ? 0 : 1;
+      Kokkos::parallel_for(
+          "compact_canonical:fill",
+          fill_teams.set_scratch_size(scratch_level,
+                                      Kokkos::PerTeam(scratch_bytes)),
+          KOKKOS_LAMBDA(const member_type& team) {
+            const int i = d_ilist(team.league_rank());
+            const int mi = loc2model(i);
+            if (mi < 0) {
+              return;
+            }
+            key_scratch lane_key(team.team_scratch(scratch_level),
+                                 neighbor_lanes);
+            key_scratch keys(team.team_scratch(scratch_level), key_capacity);
+
+            const double xi = x(i, 0);
+            const double yi = x(i, 1);
+            const double zi = x(i, 2);
+            const int jnum = d_numneigh(i);
+            const int lane = team.team_rank();
+            const std::int64_t edge = destination_row_ptr(mi);
+            constexpr unsigned long long kRejected = ~0ull;
+
+            // === Collect the sort keys of the survivors ===
+            int count = 0;
+            for (int base = 0; base < jnum; base += neighbor_lanes) {
+              const int jj = base + lane;
+              unsigned long long key = kRejected;
+              if (jj < jnum) {
+                const int j = d_neighbors(i, jj) & NEIGHMASK;
+                if (candidate_to_model(j) >= 0) {
+                  const double dx = x(j, 0) - xi;
+                  const double dy = x(j, 1) - yi;
+                  const double dz = x(j, 2) - zi;
+                  const double length2 = dx * dx + dy * dy + dz * dz;
+                  if (length2 < cutsq) {
+                    key = (static_cast(length2 * key_scale)
+                           << 48) |
+                          (static_cast(
+                               static_cast(tag(j)))
+                           << 16) |
+                          static_cast(jj);
                   }
-                  offset += target >= 0 ? 1 : 0;
-                },
-                kept);
-            edge += kept;
-            team.team_barrier();
-          }
-        });
+                }
+              }
+              lane_key(lane) = key;
+              team.team_barrier();
+              int kept = 0;
+              Kokkos::parallel_scan(
+                  Kokkos::TeamThreadRange(team, neighbor_lanes),
+                  [&](const int slot, int& offset, const bool final) {
+                    const unsigned long long value = lane_key(slot);
+                    if (final && value != kRejected) {
+                      keys(count + offset) = value;
+                    }
+                    offset += value != kRejected ? 1 : 0;
+                  },
+                  kept);
+              count += kept;
+              team.team_barrier();
+            }
+            Kokkos::Experimental::sort_team(
+                team, Kokkos::subview(keys, Kokkos::pair(0, count)));
+
+            // === Emit the row in key order ===
+            for (int slot = lane; slot < count; slot += neighbor_lanes) {
+              const int jj = static_cast(keys(slot) & 0xffffu);
+              const int j = d_neighbors(i, jj) & NEIGHMASK;
+              const int mj = candidate_to_model(j);
+              const std::int64_t position = edge + slot;
+              source(position) = static_cast(mj);
+              edge_vec(3 * position + 0) =
+                  static_cast((x(j, 0) - xi) * inv_dist);
+              edge_vec(3 * position + 1) =
+                  static_cast((x(j, 1) - yi) * inv_dist);
+              edge_vec(3 * position + 2) =
+                  static_cast((x(j, 2) - zi) * inv_dist);
+              if (source_csr) {
+                Kokkos::atomic_fetch_add(&source_counts(mj), std::int64_t{1});
+              }
+            }
+          });
+
+      // === Guard edges past the physical end of the graph ===
+      // The single-pass fill clears the rows it leaves behind itself; the
+      // counted fill reserved the storage before it filled the rows.
+      const auto storage = static_cast(edge_capacity);
+      if (storage > edge_count) {
+        Kokkos::parallel_for(
+            "compact_canonical:guards",
+            Kokkos::RangePolicy>(
+                edge_count, storage),
+            KOKKOS_LAMBDA(const std::int64_t edge) {
+              source(edge) = std::uint32_t{0};
+              edge_vec(3 * edge + 0) = 0.0f;
+              edge_vec(3 * edge + 1) = 0.0f;
+              edge_vec(3 * edge + 2) = 0.0f;
+            });
+      }
+      Kokkos::deep_copy(d_fill_total, edge_count);
+    }
+    reserve_edges(edge_count);
+    storage_count = static_cast(edge_capacity);
+    if (!source_csr) {
+      return;
+    }
 
     // === Source-major CSR and the permutation into source order ===
+    if (d_source_order.extent(0) < edge_capacity) {
+      d_source_order = Kokkos::View(
+          "compact_canonical:source_order", edge_capacity);
+    }
     auto source_row_ptr = d_source_row_ptr;
     Kokkos::parallel_scan(
         "compact_canonical:source_scan",
@@ -396,9 +1553,9 @@ class CompactCanonicalGraphKokkos : protected Pointers {
         "compact_canonical:source_cursor",
         Kokkos::RangePolicy(0, node_count_int),
         KOKKOS_LAMBDA(const int node) {
-          source_cursor(node) =
-              static_cast(source_row_ptr(node));
+          source_cursor(node) = source_row_ptr(node);
         });
+    auto source = d_source;
     auto source_order = d_source_order;
     Kokkos::parallel_for(
         "compact_canonical:source_scatter",
@@ -406,26 +1563,104 @@ class CompactCanonicalGraphKokkos : protected Pointers {
             0, edge_count),
         KOKKOS_LAMBDA(const std::int64_t edge) {
           const auto position = Kokkos::atomic_fetch_add(
-              &source_cursor(source(edge)), std::uint32_t{1});
-          source_order(position) = static_cast(edge);
+              &source_cursor(source(edge)), std::int64_t{1});
+          source_order(position) = edge;
         });
-
-    // === Guard edges past the physical end of the graph ===
+    // The cursor places the edges of a source row in the order their threads
+    // arrive, which varies from run to run, so every row is then sorted by
+    // edge position. The source order is a function of the destination-major
+    // graph alone, and the reductions that walk it -- the force assembly of
+    // DPA1 and the source-attributed per-atom virial -- sum in the same order
+    // from run to run. A row holds distinct positions, so the sorted row is
+    // the same whichever sort produces it. On a CUDA device whose candidate
+    // lists fit 512 entries, a warp sorts each row, in its registers up to
+    // that length and in memory beyond it, which no row of such a list
+    // reaches in practice; any other configuration sorts each row with a
+    // team sort.
+    bool warp_sort = false;
+#if defined(KOKKOS_ENABLE_CUDA)
+    warp_sort =
+        std::is_same_v &&
+        longest_candidates <= 16 * neighbor_lanes;
+#endif
+    if (warp_sort) {
+      constexpr int kSortWarps = 8;
+      using sort_policy =
+          Kokkos::TeamPolicy>;
+      Kokkos::parallel_for(
+          "compact_canonical:source_order_sort",
+          sort_policy((node_count_int + kSortWarps - 1) / kSortWarps,
+                      kSortWarps * 32),
+          KOKKOS_LAMBDA(const typename sort_policy::member_type& team) {
+            compact_canonical_sort_source_rows(
+                team.team_rank(), team.league_rank(), node_count_int,
+                source_row_ptr, source_order);
+          });
+    } else {
+      using sort_policy = Kokkos::TeamPolicy;
+      constexpr int sort_lanes = 32;
+      Kokkos::parallel_for(
+          "compact_canonical:source_order_sort",
+          sort_policy(node_count_int, sort_lanes),
+          KOKKOS_LAMBDA(const typename sort_policy::member_type& team) {
+            const int node = team.league_rank();
+            const auto row = Kokkos::make_pair(source_row_ptr(node),
+                                               source_row_ptr(node + 1));
+            Kokkos::Experimental::sort_team(team,
+                                            Kokkos::subview(source_order, row));
+          });
+    }
     if (storage_count > edge_count) {
       Kokkos::parallel_for(
-          "compact_canonical:guards",
+          "compact_canonical:source_order_guards",
           Kokkos::RangePolicy>(
               edge_count, storage_count),
           KOKKOS_LAMBDA(const std::int64_t edge) {
-            source(edge) = std::uint32_t{0};
-            edge_vec(3 * edge + 0) = 0.0f;
-            edge_vec(3 * edge + 1) = 0.0f;
-            edge_vec(3 * edge + 2) = 0.0f;
-            source_order(edge) = static_cast(edge);
+            source_order(edge) = edge;
           });
     }
   }
 
+  // CUDA requires a named public enclosing function for an extended lambda.
+  template 
+  void fill_rows(const Policy& policy,
+                 NeighListKokkos* list,
+                 float cutsq,
+                 float key_scale,
+                 float vector_scale,
+                 bool source_csr) {
+    const int inum = list->inum;
+    const int tail_begin = nloc_model;
+    const int tail_end = nnode_model;
+    const auto capacity = static_cast(edge_capacity);
+    auto numneigh = list->d_numneigh;
+    auto neighbors = list->d_neighbors;
+    auto loc2model = d_loc2model;
+    auto model2loc = d_model2loc;
+    auto tag = atomKK->k_tag.template view();
+    auto type = atomKK->k_type.template view();
+    auto packed = d_packed;
+    auto type_map = d_type_map;
+    auto model_type = d_model_type;
+    auto destination_row_ptr = d_destination_row_ptr;
+    auto source = d_source;
+    auto edge_vec = d_edge_vec;
+    auto source_counts = d_source_counts;
+    auto* ticket_counter = d_fill_ticket.data();
+    auto* status = d_fill_status.data();
+    auto* fill_total = d_fill_total.data();
+    Kokkos::parallel_for(
+        "compact_canonical:fill", policy,
+        KOKKOS_LAMBDA(const typename Policy::member_type& team) {
+          compact_canonical_fill_row(
+              team.team_rank(), ticket_counter, status, inum, tail_begin,
+              tail_end, capacity, fill_total, cutsq, key_scale, vector_scale,
+              numneigh, neighbors, loc2model, model2loc, tag, packed, type,
+              type_map, model_type, destination_row_ptr, source, edge_vec,
+              source_csr, source_counts);
+        });
+  }
+
   // === Node set, valid after refresh_nodes() ===
   int nloc_model;   // real local model nodes; the energy is summed over these
   int nnode_model;  // total model nodes (== nloc_model folded; + ghost
@@ -441,33 +1676,113 @@ class CompactCanonicalGraphKokkos : protected Pointers {
   typename AT::t_int_1d d_candidate_to_model;
 
   // === Compact canonical artifact, valid after build() ===
-  // The edge arrays are padded to ``storage_count`` rows: the traced program
-  // declares its edge axis with a lower bound of two, so a graph with fewer
-  // physical edges is completed with guard rows that carry a zero bond vector
-  // on node zero.
+  // The edge arrays are handed over at their allocated extent,
+  // ``storage_count`` rows, rather than at the edge count: the edge axis keeps
+  // its length while the count fluctuates from step to step, so an unchanged
+  // evaluation can be replayed, and it always meets the lower bound of two
+  // that the traced program declares. Rows past the physical end of the graph
+  // are guards that carry a zero bond vector on node zero.
   std::int64_t storage_count;
   Kokkos::View d_model_type;  // (nnode_model)
-  Kokkos::View d_source;     // (storage_count)
-  Kokkos::View d_edge_vec;           // (3 * storage_count)
+  // Source model nodes take 32 bits, since the local and ghost atoms of a rank
+  // stay below 2^31; edge positions and row pointers are 64-bit.
+  Kokkos::View d_source;  // (storage_count)
+  Kokkos::View d_edge_vec;        // (3 * storage_count)
   Kokkos::View d_destination_row_ptr;
+  // Source-major CSR, current only after a build() that asked for it.
   Kokkos::View d_source_row_ptr;
-  Kokkos::View d_source_order;
+  Kokkos::View d_source_order;
 
  private:
   // Shortest edge axis the traced program accepts.
   static constexpr std::int64_t min_storage_edges = 2;
+  // Bytes of edge storage per edge: one source index and one bond vector.
+  static constexpr std::uint64_t edge_storage_bytes =
+      sizeof(std::uint32_t) + 3 * sizeof(float);
+  // Largest edge storage sized by the candidate total rather than by the
+  // edge count read back after the fill: 1/64 of the device memory. The
+  // readback holds the host until the fill ends, so the device idles while
+  // the host launches the model, about 25 us per step on compute capability
+  // 12.0, which is 1% of the step of a system of 150,000 atoms of the
+  // smallest grade at a neighbor skin of 1 Å; the candidate storage of such a
+  // system takes 0.7 GB, 0.3 GB more than storage sized by its edges.
+  std::uint64_t bounded_storage_bytes() const {
+#if defined(KOKKOS_ENABLE_CUDA)
+    if constexpr (std::is_same_v) {
+      return typename DeviceType::execution_space()
+                 .cuda_device_prop()
+                 .totalGlobalMem /
+             64;
+    }
+#endif
+    return 0;
+  }
+
+  // Read the physical edge count of the graph just filled. The copy goes
+  // through pinned memory on the execution space instance, so the host waits
+  // for the graph build alone instead of fencing the whole device.
+  std::int64_t read_edge_count(const int node_count) {
+    const typename DeviceType::execution_space instance;
+    Kokkos::deep_copy(instance, h_edge_count,
+                      Kokkos::subview(d_destination_row_ptr, node_count));
+    instance.fence("compact_canonical:edge_count");
+    return h_edge_count();
+  }
+
+  // Grow the edge arrays to hold ``edge_count`` edges; storage never shrinks.
+  void reserve_edges(const std::int64_t edge_count) {
+    const std::int64_t storage =
+        std::max(edge_count, min_storage_edges);
+    // The vector allocation is the largest individual buffer per edge.
+    constexpr std::size_t max_storage =
+        std::numeric_limits::max() / (3 * sizeof(float));
+    if (static_cast(storage) > max_storage) {
+      error->one(FLERR,
+                 "Compact canonical graph exceeds addressable edge storage");
+    }
+    const std::size_t required = static_cast(storage);
+    if (edge_capacity < required) {
+      // A reserve of 2% covers thermal fluctuations without retaining large
+      // unused buffers, bounded by the addressable size of its arrays.
+      const std::size_t slack = required / 50 + 64;
+      edge_capacity = std::min(required + slack, max_storage);
+      d_source = Kokkos::View(
+          "compact_canonical:source", edge_capacity);
+      d_edge_vec = Kokkos::View(
+          "compact_canonical:edge_vec", edge_capacity * 3);
+      // Fresh storage is zero, so no row of it needs clearing.
+      Kokkos::deep_copy(d_fill_total, std::int64_t{0});
+    }
+  }
 
   ExecutionSpace execution_space;
+  Kokkos::View h_edge_count;
   std::vector host_type_map;
   bool extended_nodes;
   std::size_t edge_capacity;
+  // Longest candidate list of the current neighbor list, or -1 before the
+  // first build, and the candidates of all model centers.
+  int longest_candidates = -1;
+  std::int64_t candidate_total = 0;
+  // State of the single-pass fill: one status word per block of centers, in
+  // an array sized by the center count, and the ticket counter. The pack
+  // kernel clears both before every fill attempt.
+  Kokkos::View d_fill_status;
+  Kokkos::View d_fill_ticket;
+  // Fixed-point coordinates and candidate model node of every atom, one
+  // 16-byte record per atom consumed by the fill.
+  Kokkos::View d_packed;
+  // Rows of the edge storage the last build wrote, by either fill; storage
+  // past it is zero.
+  Kokkos::View d_fill_total;
 
   DAT::tdual_int_1d k_loc2model;
   DAT::tdual_int_1d k_candidate_to_model;
-  // Per-node counts are bounded by the LAMMPS neighbor limit. CSR offsets
-  // remain int64 so the compact graph retains its full global edge range.
-  Kokkos::View d_source_counts;
-  Kokkos::View d_source_cursor;
+  // Source counts and scatter cursors cover the full global edge range,
+  // including source rows assembled from many destination rows.
+  Kokkos::View d_source_counts;
+  Kokkos::View d_source_cursor;
 };
 
 }  // namespace LAMMPS_NS
diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp
index a03482a90a..cc7c6b1744 100644
--- a/source/lmp/pair_deepmd.cpp
+++ b/source/lmp/pair_deepmd.cpp
@@ -178,6 +178,11 @@ double PairDeepMD::eval_energy_with_fparam(
   dbox[7] = domain->h[3] / dist_unit_cvt_factor;  // zy
   dbox[6] = domain->h[4] / dist_unit_cvt_factor;  // zx
   dbox[3] = domain->h[5] / dist_unit_cvt_factor;  // yx
+  // A long-range module reads the cell of a periodic system only; the empty
+  // box marks a non-periodic one.
+  if (long_range && domain->nonperiodic) {
+    dbox.clear();
+  }
 
   // get coord
   for (int ii = 0; ii < nall; ++ii) {
@@ -297,6 +302,11 @@ void PairDeepMD::compute(int eflag, int vflag) {
   dbox[7] = domain->h[3] / dist_unit_cvt_factor;  // zy
   dbox[6] = domain->h[4] / dist_unit_cvt_factor;  // zx
   dbox[3] = domain->h[5] / dist_unit_cvt_factor;  // yx
+  // A long-range module reads the cell of a periodic system only; the empty
+  // box marks a non-periodic one.
+  if (long_range && domain->nonperiodic) {
+    dbox.clear();
+  }
 
   // get coord
   for (int ii = 0; ii < nall; ++ii) {
@@ -966,6 +976,36 @@ void PairDeepMD::settings(int narg, char** arg) {
   all_force.resize(numb_models);
 }
 
+void PairDeepMD::init_style() {
+  PairDeepBaseModel::init_style();
+  // A long-range module sums over every pair of atoms of a non-periodic
+  // system, or over the images of the cell of a system periodic in every
+  // dimension; the atoms of the other ranks reach it through the
+  // communicator. The model checks which of the two its module serves.
+  long_range = numb_models > 1 ? deep_pot_model_devi.has_long_range()
+                               : deep_pot.has_long_range();
+  const int periodic_dimensions =
+      domain->xperiodic + domain->yperiodic + domain->zperiodic;
+  if (long_range && periodic_dimensions != 0 && periodic_dimensions != 3) {
+    error->all(FLERR,
+               "pair style deepmd: a model with a long-range module evaluates "
+               "a system periodic in every dimension or in none; use p in "
+               "every dimension, or f, s or m in every dimension.");
+  }
+  // A non-periodic system reaches an ensemble that holds a long-range module
+  // without a cell, which only models with such a module accept. The steps
+  // without a model deviation evaluate the first model alone; the model
+  // deviation checks every model at the steps that evaluate them all.
+  if (numb_models > 1 && long_range && periodic_dimensions == 0 &&
+      !deep_pot.has_long_range()) {
+    error->all(FLERR,
+               "pair style deepmd: an ensemble that holds a model with a "
+               "long-range module evaluates a non-periodic system only when "
+               "every model carries one; use models with long-range modules "
+               "throughout, or a periodic system.");
+  }
+}
+
 /* ----------------------------------------------------------------------
    set coeffs for one or more type pairs
 ------------------------------------------------------------------------- */
diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h
index c7c40d9b18..69466f740a 100644
--- a/source/lmp/pair_deepmd.h
+++ b/source/lmp/pair_deepmd.h
@@ -46,6 +46,7 @@ class PairDeepMD : public PairDeepBaseModel {
   ~PairDeepMD() override;
   void settings(int, char**) override;
   void coeff(int, char**) override;
+  void init_style() override;
   void compute(int, int) override;
   int pack_reverse_comm(int, int, double*) override;
   void unpack_reverse_comm(int, int*, double*) override;
@@ -58,6 +59,9 @@ class PairDeepMD : public PairDeepBaseModel {
   // geometry fields are unused) for the device-resident message-passing path,
   // where ghost features are exchanged across ranks inside the forward pass.
   deepmd_compat::InputNlist make_comm_nlist();
+  // The model carries a long-range module, which reads the cell of a
+  // periodic system and no cell for a non-periodic one.
+  bool long_range = false;
 
  private:
   CommBrickDeepMD* commdata_;
diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp
index f54873dbe1..4fe527b821 100644
--- a/source/lmp/pair_deepmd_kokkos.cpp
+++ b/source/lmp/pair_deepmd_kokkos.cpp
@@ -2,6 +2,7 @@
 #ifdef LMP_KOKKOS
 #include "pair_deepmd_kokkos.h"
 
+#include 
 #include 
 #include 
 #include 
@@ -34,6 +35,7 @@ PairDeepMDKokkos::PairDeepMDKokkos(LAMMPS* lmp)
       edge_capacity(0),
       edge_vec_fp32(false),
       canonical_graph(false),
+      canonical_source_csr(false),
       device_path_ok(false),
       reverse_virial(false),
       reverse_used_host(false) {
@@ -177,6 +179,13 @@ void PairDeepMDKokkos::init_style() {
   if (std::is_same::value) {
     error->all(FLERR, "pair style deepmd/kk runs on the GPU backend only.");
   }
+  // The constructor runs before LAMMPS confirms that KOKKOS is enabled, so
+  // the virial of the graph is allocated here, once: its address stays fixed
+  // from run to run, as the replay of a graph without halo nodes keys on it.
+  if (!d_graph_virial.is_allocated()) {
+    d_graph_virial =
+        Kokkos::View("deepmd/kk:graph_virial");
+  }
   device_path_ok = deep_pot.supports_device_edge_inference();
   if (!device_path_ok) {
     error->all(
@@ -212,6 +221,17 @@ void PairDeepMDKokkos::init_style() {
   request->enable_full();
   edge_vec_fp32 = deep_pot.uses_fp32_edge_vectors();
   canonical_graph = deep_pot.uses_canonical_graph_inference();
+  canonical_source_csr = deep_pot.canonical_graph_requires_source_csr();
+  // PairDeepMD::init_style() has confined a model with a long-range module to
+  // a system periodic in every dimension or in none. The device edge path
+  // passes such a module no cell, so a periodic system reaches it through the
+  // compact canonical lower alone.
+  if (long_range && !canonical_graph && !domain->nonperiodic) {
+    error->all(FLERR,
+               "pair style deepmd/kk evaluates a periodic system with a "
+               "long-range module only through a compressed DPA4C model frozen "
+               "for a CUDA device; use pair_style deepmd for this model.");
+  }
   // Force exchange transfers three values per atom. Centroid per-atom virial
   // uses nine values even though the Kokkos full-list request runs newton off,
   // so comm_reverse_off reserves the classic host buffer for the wider mode.
@@ -387,25 +407,43 @@ int PairDeepMDKokkos::build_edges_device() {
           }
         }
       });
+  return nedge;
+}
 
-  // Compacted coordinates in model-node order for edge-input models (the graph
-  // lower ignores coordinates); only needed when virtual atoms compact them.
-  if (compact_graph.has_null_types) {
-    if ((int)d_coord_model.extent(0) < 3 * nnode_model) {
-      d_coord_model = Kokkos::View("deepmd/kk:coord_model",
-                                                        3 * nnode_model);
-    }
-    auto coord_model = d_coord_model;
-    Kokkos::parallel_for(
-        "deepmd/kk:coord", Kokkos::RangePolicy(0, nnode_model),
-        KOKKOS_LAMBDA(const int m) {
-          const int i = model2loc(m);
-          coord_model(3 * m + 0) = x(i, 0);
-          coord_model(3 * m + 1) = x(i, 1);
-          coord_model(3 * m + 2) = x(i, 2);
-        });
+template 
+const double* PairDeepMDKokkos::model_node_positions() {
+  atomKK->sync(execution_space, X_MASK);
+  auto x = atomKK->k_x.template view();
+  // A long-range module reads the positions relative to the origin of the
+  // box, as the plain pair style passes them, since its per-atom virial in a
+  // periodic cell depends on the origin; other models read them only through
+  // quantities that do not. Without virtual atoms the model nodes are the
+  // leading atoms, whose coordinates are already those positions in the
+  // length unit of the model at a conversion factor of one.
+  const double x0 = long_range ? domain->boxlo[0] : 0.0;
+  const double y0 = long_range ? domain->boxlo[1] : 0.0;
+  const double z0 = long_range ? domain->boxlo[2] : 0.0;
+  if (!compact_graph.has_null_types && dist_unit_cvt_factor == 1.0 &&
+      x0 == 0.0 && y0 == 0.0 && z0 == 0.0) {
+    return x.data();
   }
-  return nedge;
+  const int nnode_model = compact_graph.nnode_model;
+  if ((int)d_coord_model.extent(0) < 3 * nnode_model) {
+    d_coord_model = Kokkos::View("deepmd/kk:coord_model",
+                                                      3 * nnode_model);
+  }
+  auto coord_model = d_coord_model;
+  auto model2loc = compact_graph.d_model2loc;
+  const double inv_dist = 1.0 / dist_unit_cvt_factor;
+  Kokkos::parallel_for(
+      "deepmd/kk:coord", Kokkos::RangePolicy(0, nnode_model),
+      KOKKOS_LAMBDA(const int m) {
+        const int i = model2loc(m);
+        coord_model(3 * m + 0) = (x(i, 0) - x0) * inv_dist;
+        coord_model(3 * m + 1) = (x(i, 1) - y0) * inv_dist;
+        coord_model(3 * m + 2) = (x(i, 2) - z0) * inv_dist;
+      });
+  return coord_model.data();
 }
 
 template 
@@ -432,8 +470,11 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) {
     d_eatom = k_eatom.template view();
   }
   int nedge = 0;
+  // The source CSR feeds the source-attributed per-atom virial, and every
+  // force assembly of an artifact that groups edges by source.
+  const bool source_csr = canonical_source_csr || cvflag_atom;
   if (canonical_graph) {
-    compact_graph.build(list, cutoff, dist_unit_cvt_factor);
+    compact_graph.build(list, cutoff, dist_unit_cvt_factor, source_csr);
   } else {
     nedge = build_edges_device();
   }
@@ -441,32 +482,39 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) {
   const int nnode_m = compact_graph.nnode_model;  // all nodes (force / virial)
 
   // Energy is per local node; force / virial span the model node set, which is
-  // the local atoms (folded) or local + real ghost atoms (extended, up to
-  // nall). The two extents grow independently: under domain decomposition
-  // ``nlocal`` and ``nall`` need not move together, so a shared guard could
-  // leave the energy buffer short when ``nlocal`` grows while ``nall`` does
-  // not.
+  // the local atoms (folded) or local + real ghost atoms (extended). The two
+  // extents grow independently: under domain decomposition ``nlocal`` and the
+  // node count need not move together, so a shared guard could leave the
+  // energy buffer short when ``nlocal`` grows while the node count does not.
+  // Sized by the node set rather than by every atom, the buffers of a folded
+  // run keep their addresses while the ghost count fluctuates, which lets an
+  // unchanged evaluation be replayed.
   const int nall = atom->nlocal + atom->nghost;
   if ((int)d_atom_energy.extent(0) < nlocal) {
     d_atom_energy =
         Kokkos::View("deepmd/kk:atom_energy", nlocal);
   }
-  if ((int)d_out_force.extent(0) < 3 * nall) {
+  if ((int)d_out_force.extent(0) < 3 * nnode_m) {
     d_out_force =
-        Kokkos::View("deepmd/kk:out_force", 3 * nall);
+        Kokkos::View("deepmd/kk:out_force", 3 * nnode_m);
     d_atom_virial =
-        Kokkos::View("deepmd/kk:atom_virial", 9 * nall);
+        Kokkos::View("deepmd/kk:atom_virial", 9 * nnode_m);
   }
   if (cvflag_atom && multi_rank && (int)k_reverse_virial.extent(0) < 9 * nall) {
     k_reverse_virial =
         DAT::tdual_double_1d("deepmd/kk:reverse_virial", 9 * nall);
   }
-  Kokkos::deep_copy(d_out_force, 0.0);
-  Kokkos::deep_copy(d_atom_energy, 0.0);
-  Kokkos::deep_copy(d_atom_virial, 0.0);
-
-  atomKK->sync(execution_space, X_MASK | TYPE_MASK);
-  auto x = atomKK->k_x.template view();
+  // The canonical evaluation overwrites every output entry the step reads,
+  // so only its absence leaves outputs to clear; the virial of the graph is
+  // read on the canonical path alone.
+  if (!canonical_graph || nnode_m == 0) {
+    Kokkos::deep_copy(d_out_force, 0.0);
+    Kokkos::deep_copy(d_atom_energy, 0.0);
+    Kokkos::deep_copy(d_atom_virial, 0.0);
+  }
+  if (canonical_graph && nnode_m == 0) {
+    Kokkos::deep_copy(d_graph_virial, 0.0);
+  }
 
   // Runtime frame (fparam) and per-atom (aparam) parameters, built per step
   // from the same sources as the standalone pair (compute / fix / ttm or a
@@ -520,30 +568,54 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) {
     comm_ptr = &comm_list;
   }
 
-  if ((canonical_graph && nnode_m > 0) ||
+  // A model with a long-range module evaluates the system over all the ranks
+  // that hold it, so on several ranks every rank calls it at every step, one
+  // without nodes included, and passes it the communicator.
+  void* const world_ptr =
+      long_range && multi_rank ? static_cast(&world) : nullptr;
+  if ((canonical_graph && (nnode_m > 0 || world_ptr != nullptr)) ||
       (!canonical_graph && (nloc_m > 0 || comm_ptr != nullptr))) {
     // Fully device-resident inference: raw device pointers in and out. The
-    // edge buffers are produced on the Kokkos stream and consumed by the model
-    // on PyTorch's stream, and the outputs flow back to the Kokkos scatter, so
-    // the two runtimes are bracketed by explicit synchronization: fence the
-    // Kokkos work before the model reads the edges, and synchronize the device
-    // after so the scatter sees the finished model outputs.
-    Kokkos::fence();
-    // Coordinates are model-node order: the compacted buffer when virtual atoms
-    // are present, else the local coordinates directly (the graph lower ignores
-    // them; edge-input models consume them).
+    // compact canonical model is enqueued on the Kokkos stream that builds its
+    // graph and scatters its forces, so the three stay ordered without a host
+    // synchronization. The edge-input models, and every model of an execution
+    // space without a CUDA stream to share, run on PyTorch's own stream, so
+    // the Kokkos work is fenced before the model reads the edges, and the
+    // model synchronizes the device before returning its outputs.
+    // Every edge-input model takes the positions of the model nodes, and a
+    // compact canonical model reads them only with a long-range module. Their
+    // gather is enqueued before the fence, and on the canonical path on the
+    // stream that the model shares.
     const double* coord_ptr =
-        compact_graph.has_null_types ? d_coord_model.data() : x.data();
+        (!canonical_graph || long_range) ? model_node_positions() : nullptr;
+    // A long-range module reads the cell of a periodic system, row-major with
+    // the lattice vectors as rows, in the length unit of the model, and no
+    // cell for a non-periodic one.
+    const double* h = domain->h;
+    std::array cell = {h[0], 0.0,  0.0,  h[5], h[1],
+                                  0.0,  h[4], h[3], h[2]};
+    for (double& entry : cell) {
+      entry /= dist_unit_cvt_factor;
+    }
+    const double* cell_ptr =
+        long_range && !domain->nonperiodic ? cell.data() : nullptr;
+    void* const model_stream =
+        compact_canonical_stream(typename DeviceType::execution_space());
+    if (!canonical_graph || model_stream == nullptr) {
+      Kokkos::fence();
+    }
     try {
       if (canonical_graph) {
         deep_pot.compute_canonical_graph_gpu(
-            d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(),
+            d_atom_energy.data(), d_out_force.data(), d_graph_virial.data(),
+            cvflag_atom ? d_atom_virial.data() : nullptr,
             compact_graph.d_model_type.data(), compact_graph.d_source.data(),
-            compact_graph.d_edge_vec.data(),
+            DP_CANONICAL_INDEX32, compact_graph.d_edge_vec.data(),
             compact_graph.d_destination_row_ptr.data(),
-            compact_graph.d_source_row_ptr.data(),
-            compact_graph.d_source_order.data(), nloc_m, nnode_m,
-            compact_graph.storage_count);
+            source_csr ? compact_graph.d_source_row_ptr.data() : nullptr,
+            source_csr ? compact_graph.d_source_order.data() : nullptr,
+            coord_ptr, world_ptr, cell_ptr, nloc_m, nnode_m,
+            compact_graph.storage_count, model_stream);
       } else if (edge_vec_fp32) {
         deep_pot.compute_edges_gpu(d_atom_energy.data(), d_out_force.data(),
                                    d_atom_virial.data(), coord_ptr,
@@ -569,9 +641,11 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) {
   // forces are written to the ghost slots and folded onto their owners by the
   // reverse communication that the KOKKOS package (which forces 'newton off'
   // with a full list) would otherwise skip.
-  // The scatter remains device-resident. If LAMMPS selects classic host
-  // communication, the host pack/unpack methods synchronize the force DualView
-  // and the completed fold is copied back once after all communication stages.
+  // The scatter and the fold stay device-resident: the reverse communication
+  // packs and unpacks on the device (CompactCanonicalGhostFold). Only a
+  // communication style other than the Kokkos brick communication calls the
+  // host pack/unpack methods, which synchronize the force DualView, and the
+  // completed fold is then copied back once.
   auto model2loc = compact_graph.d_model2loc;
   const double fscale = scale[1][1] * force_unit_cvt_factor;
   reverse_virial = false;
@@ -604,7 +678,7 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) {
       });
   atomKK->modified(execution_space, F_MASK);
   if (multi_rank) {
-    comm->reverse_comm(this, 3);
+    ghost_fold.fold(comm, world, this, 3);
     if (reverse_used_host) {
       atomKK->sync(execution_space, F_MASK);
     }
@@ -621,22 +695,31 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) {
   }
 
   if (vflag_global) {
-    // Sum the per-node 9-component virial and map to the LAMMPS global 6
-    // (xx, yy, zz, xy, xz, yz), matching the standalone pair's index map. The
-    // sum spans all nodes (local + extended ghost) so the reduction equals the
-    // model's reduced virial for this rank's local-centered edges.
-    auto atom_virial = d_atom_virial;
+    // Map the 9-component virial to the LAMMPS global 6 (xx, yy, zz, xy, xz,
+    // yz), matching the standalone pair's index map. The canonical artifact
+    // returns it for the whole graph; otherwise the per-node virial is summed
+    // over all nodes (local + extended ghost), which equals the model's
+    // reduced virial for this rank's local-centered edges.
     const int comp[6] = {0, 4, 8, 3, 6, 7};
-    for (int k = 0; k < 6; ++k) {
-      const int off = comp[k];
-      double vsum = 0.0;
-      Kokkos::parallel_reduce(
-          "deepmd/kk:vsum", Kokkos::RangePolicy(0, nnode_m),
-          KOKKOS_LAMBDA(const int m, double& acc) {
-            acc += atom_virial(9 * m + off);
-          },
-          vsum);
-      virial[k] += scale[1][1] * vsum * ener_unit_cvt_factor;
+    if (canonical_graph) {
+      auto h_virial = Kokkos::create_mirror_view(d_graph_virial);
+      Kokkos::deep_copy(h_virial, d_graph_virial);
+      for (int k = 0; k < 6; ++k) {
+        virial[k] += scale[1][1] * h_virial(comp[k]) * ener_unit_cvt_factor;
+      }
+    } else {
+      auto atom_virial = d_atom_virial;
+      for (int k = 0; k < 6; ++k) {
+        const int off = comp[k];
+        double vsum = 0.0;
+        Kokkos::parallel_reduce(
+            "deepmd/kk:vsum", Kokkos::RangePolicy(0, nnode_m),
+            KOKKOS_LAMBDA(const int m, double& acc) {
+              acc += atom_virial(9 * m + off);
+            },
+            vsum);
+        virial[k] += scale[1][1] * vsum * ener_unit_cvt_factor;
+      }
     }
   }
 
@@ -681,7 +764,7 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) {
         }
       }
       k_reverse_virial.template sync();
-      comm->reverse_comm(this, 9);
+      ghost_fold.fold(comm, world, this, 9);
       k_reverse_virial.sync_host();
       for (int i = 0; i < nlocal; ++i) {
         for (int k = 0; k < 9; ++k) {
diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h
index bd3689aedf..22e3410132 100644
--- a/source/lmp/pair_deepmd_kokkos.h
+++ b/source/lmp/pair_deepmd_kokkos.h
@@ -67,9 +67,11 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase {
   // Fold extended (ghost) node outputs onto their owners. The KOKKOS package
   // forces 'newton off' with a full neighbor list, disabling the integrator's
   // automatic reverse communication, so the extended multi-domain path drives
-  // it explicitly for force and centroid per-atom virial. The Kokkos overrides
-  // run device-resident with GPU-aware MPI; the plain overrides serve the
-  // host-staged path.
+  // it explicitly for force and centroid per-atom virial. The Kokkos
+  // overrides pack and unpack on the device, in the device path of the KOKKOS
+  // package and in the fold that replaces its host path
+  // (CompactCanonicalGhostFold); the plain overrides serve communication styles
+  // other than the Kokkos brick communication.
   int pack_reverse_comm(int, int, double*) override;
   void unpack_reverse_comm(int, int*, double*) override;
   int pack_reverse_comm_kokkos(int, int, DeepMDKokkosCommBuffer&) override;
@@ -81,6 +83,10 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase {
   // neighbor list, returning the edge count. Public because it launches
   // extended device lambdas, which CUDA forbids inside non-public members.
   int build_edges_device();
+  // Positions of the model nodes in model-node order and in the length unit
+  // of the model, relative to the origin of the box for a model with a
+  // long-range module; public for the same reason.
+  const double* model_node_positions();
 
  protected:
   // Model node set and, for a compact canonical artifact, the graph itself.
@@ -89,7 +95,7 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase {
       d_model_type;  // (nnode_model) edge-input type per model node
   bool multi_rank;   // domain-decomposed run -> extended (local+ghost) node set
   Kokkos::View
-      d_coord_model;  // (3 * nnode_model), NULL case
+      d_coord_model;  // (3 * nnode_model), model-unit positions of nodes
 
   // Compact edge schema: edge_index is [2 * nedge] (src rows then dst rows),
   // edge_vec is [3 * nedge]; offsets is the per-atom exclusive edge prefix.
@@ -102,8 +108,11 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase {
   // Model outputs on the device. Energy is per local atom; force and virial
   // span the model node set (up to ``nall`` under domain decomposition).
   Kokkos::View d_atom_energy;  // (nlocal)
-  Kokkos::View d_out_force;    // (3 * nall)
-  Kokkos::View d_atom_virial;  // (9 * nall)
+  Kokkos::View d_out_force;    // (3 * nnode_model)
+  Kokkos::View d_atom_virial;  // (9 * nnode_model)
+  // Virial of the whole compact canonical graph, which that artifact returns
+  // alongside the per-node outputs.
+  Kokkos::View d_graph_virial;
   DAT::tdual_double_1d
       k_reverse_virial;  // (9 * nall), atom-order ghost contributions
 
@@ -118,12 +127,17 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase {
   typename AT::t_kkacc_1d d_eatom;
 #endif
 
-  int edge_capacity;       // allocated edges in d_edge_index / d_edge_vec
-  bool edge_vec_fp32;      // model graph ABI consumes edge vectors in fp32
-  bool canonical_graph;    // compact source-only graph artifact
+  int edge_capacity;     // allocated edges in d_edge_index / d_edge_vec
+  bool edge_vec_fp32;    // model graph ABI consumes edge vectors in fp32
+  bool canonical_graph;  // compact source-only graph artifact
+  // The canonical artifact reads the source CSR on every call, not only for
+  // the per-atom virial.
+  bool canonical_source_csr;
   bool device_path_ok;     // resolved once in init_style
   bool reverse_virial;     // reverse communication operates on centroid virial
-  bool reverse_used_host;  // force reverse communication selected host staging
+  bool reverse_used_host;  // the host pack/unpack overrides ran
+  // Device-packed fold that replaces the host reverse communication.
+  CompactCanonicalGhostFold ghost_fold;
 };
 
 }  // namespace LAMMPS_NS
diff --git a/source/lmp/pair_dpa4spin.cpp b/source/lmp/pair_dpa4spin.cpp
index 0a8d6e1ba8..01ba78d4d2 100644
--- a/source/lmp/pair_dpa4spin.cpp
+++ b/source/lmp/pair_dpa4spin.cpp
@@ -107,9 +107,7 @@ static const char cite_user_deepmd_package[] =
     "}\n\n";
 
 namespace {
-// Reduced Planck constant in eV.ps. The model reports the magnetic force as
-// the energy gradient with respect to the magnetic moment, while LAMMPS stores
-// the precession force; the two differ by the factor hbar / |m|.
+// Reduced Planck constant in eV.ps, used by the default precession convention.
 constexpr double kHBar = 6.5821191e-04;
 // Positions of the LAMMPS global virial components (xx, yy, zz, xy, xz, yz)
 // within the nine-component tensor the model reports.
@@ -120,7 +118,11 @@ constexpr int kCentroidVirialMap[9] = {0, 4, 8, 3, 6, 7, 1, 2, 5};
 }  // namespace
 
 PairDPA4Spin::PairDPA4Spin(LAMMPS* lmp)
-    : Pair(lmp), scale(nullptr), cutoff(0.0), commdata_(nullptr) {
+    : Pair(lmp),
+      scale(nullptr),
+      cutoff(0.0),
+      fm_cartesian_force(0),
+      commdata_(nullptr) {
   if (lmp->citeme) {
     lmp->citeme->add(cite_user_deepmd_package);
   }
@@ -237,46 +239,81 @@ void PairDPA4Spin::settings(int narg, char** arg) {
   }
   cutoff = deep_spin.cutoff() * dist_unit_cvt_factor;
 
-  // How many values name a charge state is a property of the artifact, so the
-  // keyword can only be read once the model is loaded.
+  // The artifact determines how many values follow charge_spin.
+  // magnetic_force independently selects the convention stored in atom->fm.
   const int dim_chg_spin = deep_spin.dim_chg_spin();
   std::vector charge_spin;
+  bool magnetic_force_seen = false;
+  fm_cartesian_force = 0;
   int iarg = 1;
   while (iarg < narg) {
-    if (std::string(arg[iarg]) != "charge_spin") {
+    const std::string key(arg[iarg]);
+    if (key == "charge_spin") {
+      // One charge state holds for the whole run, so a second occurrence names
+      // no state the style could serve. Values already read mark the keyword
+      // as seen: a successful read always contributes at least one value.
+      if (!charge_spin.empty()) {
+        error->all(FLERR,
+                   "Illegal pair_style command: keyword charge_spin names the "
+                   "single charge state of the whole run, so it may appear "
+                   "only once.");
+      }
+      if (dim_chg_spin == 0) {
+        error->all(FLERR,
+                   "Illegal pair_style command: the artifact served by pair "
+                   "style " +
+                       style +
+                       " carries no charge/spin condition, so keyword "
+                       "charge_spin names nothing it can serve.");
+      }
+      if (iarg + dim_chg_spin >= narg) {
+        error->all(FLERR,
+                   "Illegal pair_style command: keyword charge_spin names a "
+                   "charge state with " +
+                       std::to_string(dim_chg_spin) + " value(s).");
+      }
+      for (int ii = 0; ii < dim_chg_spin; ++ii) {
+        charge_spin.push_back(
+            utils::numeric(FLERR, arg[iarg + 1 + ii], false, lmp));
+      }
+      iarg += 1 + dim_chg_spin;
+    } else if (key == "magnetic_force") {
+      if (magnetic_force_seen) {
+        error->all(FLERR,
+                   "Illegal pair_style command: keyword magnetic_force may "
+                   "appear only once.");
+      }
+      if (iarg + 1 >= narg) {
+        error->all(FLERR,
+                   "Illegal pair_style command: keyword magnetic_force takes "
+                   "one of precession or cartesian.");
+      }
+      const std::string value(arg[iarg + 1]);
+      if (value == "precession") {
+        fm_cartesian_force = 0;
+      } else if (value == "cartesian") {
+        fm_cartesian_force = 1;
+      } else {
+        error->all(FLERR,
+                   "Illegal pair_style command: keyword magnetic_force takes "
+                   "one of precession or cartesian, not '" +
+                       value + "'.");
+      }
+      magnetic_force_seen = true;
+      iarg += 2;
+    } else {
       error->all(FLERR, "Illegal pair_style command: pair style " + style +
                             " takes the artifact path followed by the optional "
-                            "keyword charge_spin, not '" +
-                            std::string(arg[iarg]) + "'.");
-    }
-    // One charge state holds for the whole run, so a second occurrence names
-    // no state the style could serve.  Values already read mark the keyword
-    // as seen: a successful read always contributes at least one value.
-    if (!charge_spin.empty()) {
-      error->all(FLERR,
-                 "Illegal pair_style command: keyword charge_spin names the "
-                 "single charge state of the whole run, so it may appear "
-                 "only once.");
-    }
-    if (dim_chg_spin == 0) {
-      error->all(FLERR,
-                 "Illegal pair_style command: the artifact served by pair "
-                 "style " +
-                     style +
-                     " carries no charge/spin condition, so keyword "
-                     "charge_spin names nothing it can serve.");
-    }
-    if (iarg + dim_chg_spin >= narg) {
-      error->all(FLERR,
-                 "Illegal pair_style command: keyword charge_spin names a "
-                 "charge state with " +
-                     std::to_string(dim_chg_spin) + " value(s).");
+                            "keywords charge_spin and magnetic_force, not '" +
+                            key + "'.");
     }
-    for (int ii = 0; ii < dim_chg_spin; ++ii) {
-      charge_spin.push_back(
-          utils::numeric(FLERR, arg[iarg + 1 + ii], false, lmp));
-    }
-    iarg += 1 + dim_chg_spin;
+  }
+
+  // Cartesian moments use the model's spin units and energies in eV.
+  if (fm_cartesian_force && strcmp(update->unit_style, "metal") != 0) {
+    error->all(FLERR,
+               "pair style " + style +
+                   " keyword magnetic_force cartesian requires units metal.");
   }
 
   // A charge/spin condition named on the pair_style line holds for the whole
@@ -305,6 +342,8 @@ void PairDPA4Spin::settings(int narg, char** arg) {
     }
     utils::logmesg(lmp, "  using charge_spin:  {}\n", values);
   }
+  utils::logmesg(lmp, "  magnetic_force:     {}\n",
+                 fm_cartesian_force ? "cartesian" : "precession");
 }
 
 /* ----------------------------------------------------------------------
@@ -534,18 +573,19 @@ void PairDPA4Spin::compute(int eflag, int vflag) {
   }
 
   // === Step 3. Accumulate the model outputs ===
-  // The model reports a force and a magnetic force for every atom the neighbor
-  // list covers, ghosts included; the spin atom style folds the ghost rows
-  // onto their owners in its reverse communication. Dividing the magnetic
-  // force by hbar / |m| turns the energy gradient with respect to the moment
-  // into the precession force LAMMPS stores. Multiplication by |m| preserves
-  // that relation while making a zero moment well-defined.
+  // Ghost contributions are folded onto their owners by the spin atom style.
+  // Precession scales -dE/dS by |m|/hbar; Cartesian keeps the energy-per-spin
+  // force, including at zero moment.
   const double force_scale = scale[1][1] * force_unit_cvt_factor;
-  const double magnetic_force_scale = force_scale / kHBar;
+  const double magnetic_force_scale = fm_cartesian_force
+                                          ? scale[1][1] * ener_unit_cvt_factor
+                                          : force_scale / kHBar;
   for (int ii = 0; ii < nall; ++ii) {
+    const double moment_factor = fm_cartesian_force ? 1.0 : sp[ii][3];
     for (int dd = 0; dd < 3; ++dd) {
       f[ii][dd] += dforce[3 * ii + dd] * force_scale;
-      fm[ii][dd] += dforce_mag[3 * ii + dd] * sp[ii][3] * magnetic_force_scale;
+      fm[ii][dd] +=
+          dforce_mag[3 * ii + dd] * moment_factor * magnetic_force_scale;
     }
   }
 
@@ -579,5 +619,9 @@ void* PairDPA4Spin::extract(const char* str, int& dim) {
     dim = 2;
     return (void*)scale;
   }
+  if (strcmp(str, "fm_cartesian_force") == 0) {
+    dim = 0;
+    return (void*)&fm_cartesian_force;
+  }
   return nullptr;
 }
diff --git a/source/lmp/pair_dpa4spin.h b/source/lmp/pair_dpa4spin.h
index 7d0aba8977..1f389b2142 100644
--- a/source/lmp/pair_dpa4spin.h
+++ b/source/lmp/pair_dpa4spin.h
@@ -66,8 +66,10 @@ class PairDPA4Spin : public Pair {
   PairDPA4Spin(class LAMMPS*);
   ~PairDPA4Spin() override;
 
-  // Load the artifact named by ``pair_style dpa4spin  [charge_spin
-  //  ... ]`` and fix the charge state the keyword names.
+  // Load the artifact named by ``pair_style dpa4spin 
+  // [charge_spin  ... ] [magnetic_force precession|cartesian]`` and
+  // fix the optional charge state and magnetic-force convention the keywords
+  // name. ``magnetic_force`` defaults to precession.
   void settings(int, char**) override;
   // Resolve the LAMMPS atom types onto the element list of the model.
   void coeff(int, char**) override;
@@ -98,6 +100,10 @@ class PairDPA4Spin : public Pair {
   std::vector type_idx_map;
   // LAMMPS unit system relative to the eV / Angstrom system of the model.
   double ener_unit_cvt_factor, dist_unit_cvt_factor, force_unit_cvt_factor;
+  // Convention in atom->fm: 0 = precession (|m|/hbar scaling), 1 = Cartesian
+  // (-dE/dS, including at zero moment). An int for the scalar Pair::extract
+  // API.
+  int fm_cartesian_force;
 
  private:
   CommBrickDPA4Spin* commdata_;
diff --git a/source/lmp/pair_dpa4spin_kokkos.cpp b/source/lmp/pair_dpa4spin_kokkos.cpp
index f479b27559..b328e2a3bb 100644
--- a/source/lmp/pair_dpa4spin_kokkos.cpp
+++ b/source/lmp/pair_dpa4spin_kokkos.cpp
@@ -18,9 +18,7 @@
 using namespace LAMMPS_NS;
 
 namespace {
-// Reduced Planck constant in eV.ps. The model reports the magnetic force as
-// the energy gradient with respect to the magnetic moment, while LAMMPS stores
-// the precession force; the two differ by the moment magnitude over hbar.
+// Reduced Planck constant in eV.ps, used by the default precession convention.
 constexpr double kHBar = 6.5821191e-04;
 }  // namespace
 
@@ -188,6 +186,13 @@ void PairDPA4SpinKokkos::init_style() {
   if (std::is_same::value) {
     error->all(FLERR, "pair style dpa4spin/kk runs on the GPU backend only.");
   }
+  // The constructor runs before LAMMPS confirms that KOKKOS is enabled, so
+  // the virial of the graph is allocated here, once: its address stays fixed
+  // from run to run, as the replay of a graph without halo nodes keys on it.
+  if (!d_graph_virial.is_allocated()) {
+    d_graph_virial =
+        Kokkos::View("dpa4spin/kk:graph_virial");
+  }
   // Device residency is what the compact canonical ABI exists for: it is the
   // only schema whose inputs this style can hand over without a host round
   // trip.
@@ -265,56 +270,70 @@ void PairDPA4SpinKokkos::compute(int eflag, int vflag) {
     memoryKK->create_kokkos(k_eatom, eatom, maxeatom, "dpa4spin/kk:eatom");
     d_eatom = k_eatom.template view();
   }
-  compact_graph.build(list, cutoff, dist_unit_cvt_factor);
+  // The source CSR feeds the source-attributed per-atom virial alone.
+  compact_graph.build(list, cutoff, dist_unit_cvt_factor, cvflag_atom);
   gather_moment_device();
   const int nloc_m = compact_graph.nloc_model;    // local nodes (energy)
   const int nnode_m = compact_graph.nnode_model;  // all nodes (force / virial)
 
   // Energy is per local node; force, magnetic force and virial span the model
   // node set, which is the local atoms (folded) or local + real ghost atoms
-  // (extended, up to nall). The two extents grow independently: under domain
-  // decomposition ``nlocal`` and ``nall`` need not move together, so a shared
-  // guard could leave the energy buffer short when ``nlocal`` grows while
-  // ``nall`` does not.
+  // (extended). The two extents grow independently: under domain
+  // decomposition ``nlocal`` and the node count need not move together, so a
+  // shared guard could leave the energy buffer short when ``nlocal`` grows
+  // while the node count does not. Sized by the node set rather than by every
+  // atom, the buffers of a folded run keep their addresses while the ghost
+  // count fluctuates, which lets an unchanged evaluation be replayed.
   const int nall = atom->nlocal + atom->nghost;
   if ((int)d_atom_energy.extent(0) < nlocal) {
     d_atom_energy =
         Kokkos::View("dpa4spin/kk:atom_energy", nlocal);
   }
-  if ((int)d_out_force.extent(0) < 3 * nall) {
+  if ((int)d_out_force.extent(0) < 3 * nnode_m) {
     d_out_force =
-        Kokkos::View("dpa4spin/kk:out_force", 3 * nall);
+        Kokkos::View("dpa4spin/kk:out_force", 3 * nnode_m);
     d_out_force_mag = Kokkos::View(
-        "dpa4spin/kk:out_force_mag", 3 * nall);
-    d_atom_virial =
-        Kokkos::View("dpa4spin/kk:atom_virial", 9 * nall);
+        "dpa4spin/kk:out_force_mag", 3 * nnode_m);
+    d_atom_virial = Kokkos::View("dpa4spin/kk:atom_virial",
+                                                      9 * nnode_m);
   }
   if (cvflag_atom && multi_rank && (int)k_reverse_virial.extent(0) < 9 * nall) {
     k_reverse_virial =
         DAT::tdual_double_1d("dpa4spin/kk:reverse_virial", 9 * nall);
   }
-  Kokkos::deep_copy(d_out_force, 0.0);
-  Kokkos::deep_copy(d_out_force_mag, 0.0);
-  Kokkos::deep_copy(d_atom_energy, 0.0);
-  Kokkos::deep_copy(d_atom_virial, 0.0);
+  // The evaluation overwrites every output entry the step reads, so only its
+  // absence leaves outputs to clear.
+  if (nnode_m == 0) {
+    Kokkos::deep_copy(d_out_force, 0.0);
+    Kokkos::deep_copy(d_out_force_mag, 0.0);
+    Kokkos::deep_copy(d_atom_energy, 0.0);
+    Kokkos::deep_copy(d_atom_virial, 0.0);
+    Kokkos::deep_copy(d_graph_virial, 0.0);
+  }
 
   if (nnode_m > 0) {
     // Fully device-resident inference: raw device pointers in and out. The
-    // graph and the moment are produced on the Kokkos stream and consumed by
-    // the model on PyTorch's stream, and the outputs flow back to the Kokkos
-    // scatter, so the two runtimes are bracketed by explicit synchronization:
-    // fence the Kokkos work before the model reads its inputs, and synchronize
-    // the device after so the scatter sees the finished model outputs.
-    Kokkos::fence();
+    // model is enqueued on the Kokkos stream that produces the graph and the
+    // moment and scatters the outputs, so the three stay ordered without a
+    // host synchronization. An execution space without a CUDA stream to
+    // share runs the model on PyTorch's own stream, so its Kokkos work is
+    // fenced before the model reads its inputs.
+    void* const model_stream =
+        compact_canonical_stream(typename DeviceType::execution_space());
+    if (model_stream == nullptr) {
+      Kokkos::fence();
+    }
     try {
       deep_spin.compute_canonical_graph_gpu(
           d_atom_energy.data(), d_out_force.data(), d_out_force_mag.data(),
-          d_atom_virial.data(), compact_graph.d_model_type.data(),
-          compact_graph.d_source.data(), compact_graph.d_edge_vec.data(),
+          d_graph_virial.data(), cvflag_atom ? d_atom_virial.data() : nullptr,
+          compact_graph.d_model_type.data(), compact_graph.d_source.data(),
+          DP_CANONICAL_INDEX32, compact_graph.d_edge_vec.data(),
           compact_graph.d_destination_row_ptr.data(),
-          compact_graph.d_source_row_ptr.data(),
-          compact_graph.d_source_order.data(), d_moment.data(), nloc_m, nnode_m,
-          compact_graph.storage_count);
+          cvflag_atom ? compact_graph.d_source_row_ptr.data() : nullptr,
+          cvflag_atom ? compact_graph.d_source_order.data() : nullptr,
+          d_moment.data(), nloc_m, nnode_m, compact_graph.storage_count,
+          model_stream);
     } catch (deepmd_compat::deepmd_exception& e) {
       error->one(FLERR, e.what());
     }
@@ -327,13 +346,18 @@ void PairDPA4SpinKokkos::compute(int eflag, int vflag) {
   // forces are written to the ghost slots and folded onto their owners by the
   // reverse communication that the KOKKOS package (which forces 'newton off'
   // with a full list) would otherwise skip.
-  // The scatter remains device-resident. If LAMMPS selects classic host
-  // communication, the host pack/unpack methods synchronize the force
-  // DualViews and the completed fold is copied back once after all
-  // communication stages.
+  // The scatter and the fold stay device-resident: the reverse communication
+  // packs and unpacks on the device (CompactCanonicalGhostFold). Only a
+  // communication style other than the Kokkos brick communication calls the
+  // host pack/unpack methods, which synchronize the force DualViews, and the
+  // completed fold is then copied back once.
   auto model2loc = compact_graph.d_model2loc;
   const double fscale = scale[1][1] * force_unit_cvt_factor;
-  const double fmscale = scale[1][1] * force_unit_cvt_factor / kHBar;
+  // Precession: energy gradient times |m| / hbar. Cartesian: energy gradient
+  // with only the energy-per-model-spin unit conversion (no |m|, no hbar).
+  const int cartesian = fm_cartesian_force;
+  const double fmscale =
+      cartesian ? scale[1][1] * ener_unit_cvt_factor : fscale / kHBar;
   reverse_virial = false;
   reverse_used_host = false;
   // The KOKKOS package runs 'newton off', so the integrator's force_clear only
@@ -364,19 +388,18 @@ void PairDPA4SpinKokkos::compute(int eflag, int vflag) {
       "dpa4spin/kk:scatter_f", Kokkos::RangePolicy(0, nnode_m),
       KOKKOS_LAMBDA(const int m) {
         const int i = model2loc(m);
-        // The precession force carried by ``fm`` is the energy gradient with
-        // respect to the moment, rescaled by the moment magnitude over hbar.
-        const double moment = fmscale * sp(i, 3);
+        // Precession multiplies by the moment magnitude; Cartesian does not.
+        const double moment_factor = cartesian ? 1.0 : sp(i, 3);
         f(i, 0) += fscale * out_force(3 * m + 0);
         f(i, 1) += fscale * out_force(3 * m + 1);
         f(i, 2) += fscale * out_force(3 * m + 2);
-        fm(i, 0) += moment * out_force_mag(3 * m + 0);
-        fm(i, 1) += moment * out_force_mag(3 * m + 1);
-        fm(i, 2) += moment * out_force_mag(3 * m + 2);
+        fm(i, 0) += fmscale * moment_factor * out_force_mag(3 * m + 0);
+        fm(i, 1) += fmscale * moment_factor * out_force_mag(3 * m + 1);
+        fm(i, 2) += fmscale * moment_factor * out_force_mag(3 * m + 2);
       });
   atomKK->modified(execution_space, F_MASK | FM_MASK);
   if (multi_rank) {
-    comm->reverse_comm(this, 6);
+    ghost_fold.fold(comm, world, this, 6);
     if (reverse_used_host) {
       atomKK->sync(execution_space, F_MASK | FM_MASK);
     }
@@ -393,22 +416,13 @@ void PairDPA4SpinKokkos::compute(int eflag, int vflag) {
   }
 
   if (vflag_global) {
-    // Sum the per-node 9-component virial and map to the LAMMPS global 6
-    // (xx, yy, zz, xy, xz, yz), matching the standalone pair's index map. The
-    // sum spans all nodes (local + extended ghost) so the reduction equals the
-    // model's reduced virial for this rank's local-centered edges.
-    auto atom_virial = d_atom_virial;
+    // Map the 9-component virial of the graph to the LAMMPS global 6 (xx, yy,
+    // zz, xy, xz, yz), matching the standalone pair's index map.
+    auto h_virial = Kokkos::create_mirror_view(d_graph_virial);
+    Kokkos::deep_copy(h_virial, d_graph_virial);
     const int comp[6] = {0, 4, 8, 3, 6, 7};
     for (int k = 0; k < 6; ++k) {
-      const int off = comp[k];
-      double vsum = 0.0;
-      Kokkos::parallel_reduce(
-          "dpa4spin/kk:vsum", Kokkos::RangePolicy(0, nnode_m),
-          KOKKOS_LAMBDA(const int m, double& acc) {
-            acc += atom_virial(9 * m + off);
-          },
-          vsum);
-      virial[k] += scale[1][1] * vsum * ener_unit_cvt_factor;
+      virial[k] += scale[1][1] * h_virial(comp[k]) * ener_unit_cvt_factor;
     }
   }
 
@@ -453,7 +467,7 @@ void PairDPA4SpinKokkos::compute(int eflag, int vflag) {
         }
       }
       k_reverse_virial.template sync();
-      comm->reverse_comm(this, 9);
+      ghost_fold.fold(comm, world, this, 9);
       k_reverse_virial.sync_host();
       for (int i = 0; i < nlocal; ++i) {
         for (int k = 0; k < 9; ++k) {
diff --git a/source/lmp/pair_dpa4spin_kokkos.h b/source/lmp/pair_dpa4spin_kokkos.h
index 1485e6e336..794e0c6cd9 100644
--- a/source/lmp/pair_dpa4spin_kokkos.h
+++ b/source/lmp/pair_dpa4spin_kokkos.h
@@ -38,7 +38,8 @@ using DPA4SpinKokkosCommBuffer = DAT::tdual_double_1d;
 
 // GPU-resident inference for exported native-spin ``.pt2`` models whose forward
 // consumes the compact canonical graph: a dual-CSR neighbor topology with
-// uint32 indices and float32 edge vectors, plus the per-node magnetic moment.
+// 32-bit source indices, int64 row pointers and positions, and float32 edge
+// vectors, plus the per-node magnetic moment.
 // It is dispatched through ``DeepSpin::compute_canonical_graph_gpu``.
 //
 // The neighbor list, the compact graph, the moment and the model outputs all
@@ -68,8 +69,10 @@ class PairDPA4SpinKokkos : public PairDPA4Spin, public KokkosBase {
   // forces 'newton off' with a full neighbor list, disabling the integrator's
   // automatic reverse communication, so the extended multi-domain path drives
   // it explicitly for force, magnetic force and centroid per-atom virial. The
-  // Kokkos overrides run device-resident with GPU-aware MPI; the plain
-  // overrides serve the host-staged path.
+  // Kokkos overrides pack and unpack on the device, in the device path of the
+  // KOKKOS package and in the fold that replaces its host path
+  // (CompactCanonicalGhostFold); the plain overrides serve communication styles
+  // other than the Kokkos brick communication.
   int pack_reverse_comm(int, int, double*) override;
   void unpack_reverse_comm(int, int*, double*) override;
   int pack_reverse_comm_kokkos(int, int, DPA4SpinKokkosCommBuffer&) override;
@@ -94,9 +97,11 @@ class PairDPA4SpinKokkos : public PairDPA4Spin, public KokkosBase {
   // force and virial span the model node set (up to ``nall`` under domain
   // decomposition).
   Kokkos::View d_atom_energy;    // (nlocal)
-  Kokkos::View d_out_force;      // (3 * nall)
-  Kokkos::View d_out_force_mag;  // (3 * nall)
-  Kokkos::View d_atom_virial;    // (9 * nall)
+  Kokkos::View d_out_force;      // (3 * nnode_model)
+  Kokkos::View d_out_force_mag;  // (3 * nnode_model)
+  Kokkos::View d_atom_virial;    // (9 * nnode_model)
+  // Virial of the whole compact canonical graph.
+  Kokkos::View d_graph_virial;
   DAT::tdual_double_1d
       k_reverse_virial;  // (9 * nall), atom-order ghost contributions
 
@@ -112,7 +117,9 @@ class PairDPA4SpinKokkos : public PairDPA4Spin, public KokkosBase {
 #endif
 
   bool reverse_virial;     // reverse communication operates on centroid virial
-  bool reverse_used_host;  // force reverse communication selected host staging
+  bool reverse_used_host;  // the host pack/unpack overrides ran
+  // Device-packed fold that replaces the host reverse communication.
+  CompactCanonicalGhostFold ghost_fold;
 };
 
 }  // namespace LAMMPS_NS
diff --git a/source/lmp/plugin/CMakeLists.txt b/source/lmp/plugin/CMakeLists.txt
index b87464a03d..450bcaf881 100644
--- a/source/lmp/plugin/CMakeLists.txt
+++ b/source/lmp/plugin/CMakeLists.txt
@@ -188,9 +188,13 @@ if(DEFINED LAMMPS_SOURCE_ROOT OR DEFINED LAMMPS_VERSION)
     set_target_properties(${libname} PROPERTIES LINK_FLAGS
                                                 "-Wl,-undefined,dynamic_lookup")
   else()
+    # When installed with the C/C++ interface, LAMMPS opens the plugin through
+    # its link in lib/deepmd_lmp, where $ORIGIN names that directory, so
+    # $ORIGIN/.. reaches the DeePMD-kit libraries installed beside the plugin.
     set_target_properties(
-      ${libname} PROPERTIES INSTALL_RPATH "$ORIGIN;${BACKEND_LIBRARY_PATH}"
-                            LINK_FLAGS "-rdynamic")
+      ${libname}
+      PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/..;${BACKEND_LIBRARY_PATH}"
+                 LINK_FLAGS "-rdynamic")
   endif()
   target_compile_definitions(
     ${libname}
diff --git a/source/lmp/plugin/deepmdplugin.cpp b/source/lmp/plugin/deepmdplugin.cpp
index 9bc0b685d2..46922f4ff1 100644
--- a/source/lmp/plugin/deepmdplugin.cpp
+++ b/source/lmp/plugin/deepmdplugin.cpp
@@ -13,6 +13,7 @@
 #endif
 #include "pair_deepmd.h"
 #include "pair_deepspin.h"
+#include "pair_dpa4spin.h"
 #include "version.h"
 #if LAMMPS_VERSION_NUMBER >= 20220328
 #include "pppm_dplr.h"
@@ -22,6 +23,7 @@ using namespace LAMMPS_NS;
 
 static Pair* pairdeepmd(LAMMPS* lmp) { return new PairDeepMD(lmp); }
 static Pair* pairdeepspin(LAMMPS* lmp) { return new PairDeepSpin(lmp); }
+static Pair* pairdpa4spin(LAMMPS* lmp) { return new PairDPA4Spin(lmp); }
 
 #ifdef LMP_KOKKOS
 // Runtime plugins do not consume the PairStyle declarations used by LAMMPS'
@@ -109,6 +111,15 @@ extern "C" void lammpsplugin_init(void* lmp, void* handle, void* regfunc) {
   plugin.handle = handle;
   (*register_plugin)(&plugin, lmp);
 
+  plugin.version = LAMMPS_VERSION;
+  plugin.style = "pair";
+  plugin.name = "dpa4spin";
+  plugin.info = "dpa4spin pair style " STR_GIT_SUMM;
+  plugin.author = "Tiancheng Li";
+  plugin.creator.v1 = (lammpsplugin_factory1*)&pairdpa4spin;
+  plugin.handle = handle;
+  (*register_plugin)(&plugin, lmp);
+
   plugin.style = "compute";
   plugin.name = "deeptensor/atom";
   plugin.info = "compute deeptensor/atom " STR_GIT_SUMM;
diff --git a/source/lmp/tests/lammps_test_utils.py b/source/lmp/tests/lammps_test_utils.py
index e3f33f5cb1..ff73605625 100644
--- a/source/lmp/tests/lammps_test_utils.py
+++ b/source/lmp/tests/lammps_test_utils.py
@@ -12,6 +12,7 @@
 )
 
 import os
+import signal
 import subprocess as sp
 import sys
 import tempfile
@@ -136,6 +137,7 @@ def run_mpi_pair_runner(
     runner_args: list[str] | None = None,
     output_columns: tuple[tuple[str, int], ...] = (("forces", 3), ("virials", 9)),
     capture: bool = False,
+    timeout: float | None = None,
 ) -> dict[str, Any]:
     """Invoke a DPA MPI runner and parse its energy/per-atom output.
 
@@ -144,7 +146,9 @@ def run_mpi_pair_runner(
     model-specific wrapper retains its own defaults and explanatory docstring.
 
     If ``capture`` is true, skip parsing and return the subprocess result as
-    ``{"returncode": int, "stdout": str, "stderr": str}``.
+    ``{"returncode": int, "stdout": str, "stderr": str}``. ``timeout`` bounds
+    the run in seconds, after which the launcher and its ranks are killed and
+    ``subprocess.TimeoutExpired`` is raised; ``None`` leaves it unbounded.
     """
     with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f:
         output_path = Path(f.name)
@@ -167,15 +171,30 @@ def run_mpi_pair_runner(
             argv.extend(extra_args)
         if runner_args:
             argv.extend(runner_args)
+        # The launcher and the ranks of a bounded run share a process group of
+        # their own, so that a timeout ends the ranks with it rather than
+        # leaving them holding their devices.
+        with sp.Popen(
+            argv,
+            stdout=sp.PIPE if capture else None,
+            stderr=sp.PIPE if capture else None,
+            text=True,
+            start_new_session=timeout is not None,
+        ) as process:
+            try:
+                stdout, stderr = process.communicate(timeout=timeout)
+            except sp.TimeoutExpired:
+                os.killpg(process.pid, signal.SIGKILL)
+                process.communicate()
+                raise
         if capture:
-            proc = sp.run(argv, capture_output=True, text=True)
             return {
-                "returncode": proc.returncode,
-                "stdout": proc.stdout,
-                "stderr": proc.stderr,
+                "returncode": process.returncode,
+                "stdout": stdout,
+                "stderr": stderr,
             }
-
-        sp.check_call(argv)
+        if process.returncode != 0:
+            raise sp.CalledProcessError(process.returncode, argv)
         lines = output_path.read_text().strip().splitlines()
         rows = np.array(
             [list(map(float, line.split())) for line in lines[1:]],
diff --git a/source/lmp/tests/run_mpi_dpa4spin_conventions.py b/source/lmp/tests/run_mpi_dpa4spin_conventions.py
new file mode 100644
index 0000000000..d73d3e343f
--- /dev/null
+++ b/source/lmp/tests/run_mpi_dpa4spin_conventions.py
@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Run native-spin conventions through a real host or CUDA Kokkos plugin.
+
+Each MPI process uses the same GPU, so two-rank coverage needs only one device.
+Host-staged communication deliberately does not require CUDA-aware MPI.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+import argparse
+import ctypes
+import json
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+from lammps import (
+    LMP_STYLE_GLOBAL,
+    LMP_TYPE_VECTOR,
+    lammps,
+)
+from mpi4py import (
+    MPI,
+)
+
+
+def _pair_pointer(lmp, name, pointer_type):
+    extract = lmp.lib.lammps_extract_pair
+    previous = extract.restype
+    extract.restype = pointer_type
+    try:
+        pointer = extract(lmp.lmp, name.encode())
+    finally:
+        extract.restype = previous
+    assert pointer, f"Pair::extract({name!r}) returned null"
+    return pointer
+
+
+def _snapshot(lmp, virial=False):
+    # gather() returns atom-ID order, including contributions reverse-communicated
+    # from ghost nodes. All ranks must participate in these collective calls.
+    def gather(name, count):
+        return np.array(lmp.gather(name, 1, count)).reshape(-1, count).tolist()
+
+    # The pressure compute restricted to the pair reports the global pair
+    # virial (xx, yy, zz, xy, xz, yz) over the volume in pressure units; the
+    # volume over nktv2p turns it into eV (units metal).
+    pressure = lmp.extract_compute("pvirial", LMP_STYLE_GLOBAL, LMP_TYPE_VECTOR)
+    volume = lmp.get_thermo("vol") / lmp.extract_global("nktv2p")
+    snapshot = {
+        "ids": list(lmp.gather_atoms("id", 0, 1)),
+        "force": gather("f", 3),
+        "magnetic_force": gather("c_fmprop", 3),
+        "spin": gather("c_spprop", 4),
+        "energy": lmp.get_thermo("pe"),
+        "global_virial": [volume * pressure[k] for k in range(6)],
+    }
+    if virial:
+        # centroid/stress/atom reports the negative per-atom virial in
+        # pressure*volume units; dividing by -nktv2p gives eV (units metal).
+        scale = -1.0 / lmp.extract_global("nktv2p")
+        snapshot["virial"] = (scale * np.array(gather("c_vatom", 9))).tolist()
+    return snapshot
+
+
+def _cmdargs(kokkos):
+    cmdargs = ["-log", "none", "-screen", "none"]
+    if kokkos:
+        cmdargs += [
+            "-k",
+            "on",
+            "g",
+            "1",
+            "-pk",
+            "kokkos",
+            "neigh",
+            "full",
+            "newton",
+            "off",
+            "comm",
+            "host",
+            "gpu/aware",
+            "off",
+        ]
+    return cmdargs
+
+
+def _run(lmp, args, selection, overlay=False):
+    comm = MPI.COMM_WORLD
+    assert lmp.get_mpi_comm().size == args.ranks
+    pair = "dpa4spin/kk" if args.kokkos else "dpa4spin"
+    atom = "spin/kk" if args.kokkos else "spin"
+    integrator = "verlet/kk" if args.kokkos else "verlet"
+    assert lmp.has_style("pair", pair), f"Missing plugin pair style {pair}"
+    pair_args = "" if selection == "default" else f" magnetic_force {selection}"
+    substyle = f"{pair} {args.model.resolve()}{pair_args}"
+    if overlay:
+        style = "hybrid/overlay/kk" if args.kokkos else "hybrid/overlay"
+        pair_command = f"pair_style {style} {substyle} {substyle}"
+        coefficients = [f"pair_coeff * * {pair} {i} Ni O" for i in (1, 2)]
+    else:
+        style = pair
+        pair_command = f"pair_style {substyle}"
+        coefficients = ["pair_coeff * * Ni O"]
+    for command in (
+        "units metal",
+        f"processors {comm.size} 1 1",
+        f"atom_style {atom}",
+        "atom_modify map yes",
+        f"read_data {args.data.resolve()}",
+        # No global -sf kk: select the matching integrator explicitly too.
+        f"run_style {integrator}",
+        "mass 1 58.6934",
+        "mass 2 15.9994",
+        pair_command,
+        *coefficients,
+        "neighbor 1.0 bin",
+        "compute fmprop all property/atom fmx fmy fmz",
+        "compute spprop all property/atom spx spy spz sp",
+        "compute pvirial all pressure NULL pair",
+        "run 0",
+    ):
+        lmp.command(command)
+    assert lmp.extract_pair_dimension("fm_cartesian_force") == 0
+    flag = _pair_pointer(lmp, "fm_cartesian_force", ctypes.POINTER(ctypes.c_int))[0]
+    result = {
+        "convention": flag,
+        "pair_style": lmp.extract_global("pair_style"),
+        "run_style_active": lmp.eval(f"is_active(run_style,{integrator})"),
+        "atom_style": lmp.extract_global("atom_style"),
+        "kokkos_active": lmp.extract_setting("kokkos_active"),
+        "kokkos_ngpus": lmp.extract_setting("kokkos_ngpus"),
+        "newton_pair": lmp.extract_setting("newton_pair"),
+        "local_counts": comm.allgather(lmp.extract_setting("nlocal")),
+        "ghost_counts": comm.allgather(lmp.extract_setting("nghost")),
+        "snapshots": [_snapshot(lmp)],
+    }
+    assert result["pair_style"] == style
+    assert result["atom_style"] == atom
+    # Repeated computes expose stale ghost magnetic-force accumulation.
+    for _ in range(2):
+        lmp.command("run 0")
+        result["snapshots"].append(_snapshot(lmp))
+    # A per-atom centroid virial switches the graph build and the operator to
+    # their source-attributed path; the snapshots above cover the path without
+    # it. A compute defined between runs is scheduled by every later run.
+    lmp.command("compute vatom all centroid/stress/atom NULL pair")
+    result["virial_snapshots"] = []
+    for _ in range(2):
+        lmp.command("run 0")
+        result["virial_snapshots"].append(_snapshot(lmp, virial=True))
+    if not overlay:
+        scale = _pair_pointer(
+            lmp, "scale", ctypes.POINTER(ctypes.POINTER(ctypes.c_double))
+        )
+        scale[1][1] = 2.75
+        lmp.command("run 0")
+        result["scaled"] = _snapshot(lmp, virial=True)
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("data", type=Path)
+    parser.add_argument("model", type=Path)
+    parser.add_argument("output", type=Path)
+    parser.add_argument("--kokkos", action="store_true")
+    parser.add_argument("--ranks", type=int, default=1)
+    args = parser.parse_args()
+    assert MPI.COMM_WORLD.size == args.ranks, (
+        f"Expected {args.ranks} MPI ranks, got {MPI.COMM_WORLD.size}: "
+        f"{MPI.Get_library_version()}"
+    )
+    results = {}
+    library = None
+    try:
+        for selection, overlay in (
+            ("default", False),
+            ("precession", False),
+            ("cartesian", False),
+            ("precession", True),
+            ("cartesian", True),
+        ):
+            with lammps(cmdargs=_cmdargs(args.kokkos), comm=MPI.COMM_WORLD) as lmp:
+                library = lmp.lib
+                key = f"overlay_{selection}" if overlay else selection
+                results[key] = _run(lmp, args, selection, overlay)
+        if MPI.COMM_WORLD.rank == 0:
+            args.output.write_text(json.dumps(results))
+    finally:
+        # close() destroys each instance but keeps Kokkos initialized for later
+        # instances. Finalize once, even on error, before mpi4py finalizes MPI.
+        if args.kokkos and library is not None:
+            library.lammps_kokkos_finalize()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py
index a574cef821..55d5317943 100644
--- a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py
+++ b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py
@@ -51,6 +51,13 @@
     "domain decomposition (nswap>0). Pass '1 1 1' for a single-rank "
     "reference run on the same archive (single-artifact dispatch).",
 )
+parser.add_argument(
+    "--boundary",
+    type=str,
+    default="p p p",
+    help="LAMMPS boundary style. Default 'p p p'; a model with a long-range "
+    "module takes 'f f f', or 'p p p' when its module serves periodic systems.",
+)
 parser.add_argument(
     "--pair-coeff",
     type=str,
@@ -127,9 +134,24 @@
     "to exercise sendlist composition changes from real-atom motion "
     "rather than only from NULL motion.",
 )
+parser.add_argument(
+    "--kokkos",
+    action="store_true",
+    help="Run pair_style deepmd/kk on the Kokkos CUDA device of each rank, "
+    "with a full neighbor list, newton off and host communication, as the "
+    "device path of LAMMPS + Kokkos runs it without a CUDA-aware MPI. Rank r "
+    "takes device r modulo the visible devices, as DeepPotPTExpt does.",
+)
 args = parser.parse_args()
 
-lammps = PyLammps()
+cmdargs = None
+if args.kokkos:
+    import torch
+
+    devices = min(comm.Get_size(), torch.cuda.device_count())
+    cmdargs = ["-k", "on", "g", str(devices), "-sf", "kk", "-pk", "kokkos"]
+    cmdargs += ["neigh", "full", "newton", "off", "comm", "host", "gpu/aware", "off"]
+lammps = PyLammps(cmdargs=cmdargs)
 # Force the requested domain decomposition. The default "2 1 1"
 # combined with the simulation box guarantees nswap > 0 on the C++
 # side, so DeepPotPTExpt routes to the with-comm AOTI artifact. Pass
@@ -137,7 +159,7 @@
 # (the regular artifact handles nswap==0).
 lammps.processors(args.processors)
 lammps.units("metal")
-lammps.boundary("p p p")
+lammps.boundary(args.boundary)
 lammps.atom_style("atomic")
 # ``atom_modify map yes`` is required when single-rank dispatch goes
 # through the regular artifact of a use_loc_mapping=False .pt2: the
@@ -187,7 +209,9 @@
     else:
         lammps.velocity(f"nullgroup set {args.null_vx:.6f} 0.0 0.0 units box")
 
-pair_style_line = f"deepmd {args.PB_FILE}"
+# The Kokkos style is named explicitly, so that a plugin without it fails
+# instead of falling back to the host style.
+pair_style_line = f"{'deepmd/kk' if args.kokkos else 'deepmd'} {args.PB_FILE}"
 if args.pair_extra:
     pair_style_line += f" {args.pair_extra}"
 lammps.pair_style(pair_style_line)
diff --git a/source/lmp/tests/test_dpa4spin_kokkos_magnetic_force.py b/source/lmp/tests/test_dpa4spin_kokkos_magnetic_force.py
new file mode 100644
index 0000000000..990d937be6
--- /dev/null
+++ b/source/lmp/tests/test_dpa4spin_kokkos_magnetic_force.py
@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Actual CUDA /kk convention and MPI regression, separate from host FD tests.
+
+Run source/install/test_lammps_kokkos.sh after test_cc_local.sh on a CUDA host.
+That job supplies a compressed canonical fixture and a matching Kokkos LAMMPS
+library/plugin, and sets DPA4SPIN_KOKKOS_REQUIRED=1: no missing prerequisite may
+silently skip it. This checks adapter equivalence, not GPU energy differences.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+import importlib
+import json
+import os
+import shutil
+import signal
+import subprocess
+import sys
+import zipfile
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+import pytest
+from write_lmp_data import (
+    write_lmp_data_spin,
+)
+
+HBAR_METAL = 6.5821191e-04
+_MODEL = (
+    Path(__file__).resolve().parents[2]
+    / "tests/infer/deeppot_dpa4c_spin_lmp_kokkos.pt2"
+)
+_BOX = np.array([0.0, 14.0, 0.0, 14.0, 0.0, 14.0, 0.0, 0.0, 0.0])
+# Both ranks own atoms and interact across the x=7 domain boundary. The box
+# exceeds twice the cutoff, as required by the single-rank folded /kk graph.
+_COORD = np.array(
+    [
+        [5.5, 4.0, 4.0],
+        [7.7, 4.4, 4.1],
+        [5.8, 4.8, 4.0],
+        [4.9, 4.2, 4.6],
+        [8.1, 5.0, 4.3],
+        [7.9, 3.7, 4.7],
+    ]
+)
+_SPIN = np.array(
+    [
+        [0.0, 0.0, 0.0],
+        [-0.07, 0.09, 0.03],
+        [0.02, -0.06, 0.08],
+        [0.0, 0.0, 0.0],
+        [0.0, 0.0, 0.0],
+        [0.0, 0.0, 0.0],
+    ]
+)
+_TYPE = np.array([1, 1, 1, 2, 2, 2], dtype=np.int32)
+
+
+def _unavailable(reason):
+    if os.environ.get("DPA4SPIN_KOKKOS_REQUIRED") == "1":
+        pytest.fail(reason)
+    pytest.skip(reason)
+
+
+@pytest.fixture(scope="module")
+def kokkos_model():
+    # The ordinary LAMMPS suite also runs on CUDA, but uses a non-Kokkos wheel.
+    # Only the dedicated matching library/plugin job opts into this regression.
+    if os.environ.get("DPA4SPIN_KOKKOS_REQUIRED") != "1":
+        pytest.skip("Enable the dedicated CUDA Kokkos plugin regression explicitly")
+    try:
+        torch = importlib.import_module("torch")
+    except ImportError as exc:
+        _unavailable(f"CUDA PyTorch unavailable: {exc}")
+        raise
+    if not torch.cuda.is_available():
+        _unavailable("CUDA device required for the actual dpa4spin/kk adapter")
+    for module in ("lammps", "mpi4py"):
+        if importlib.util.find_spec(module) is None:
+            _unavailable(f"Required Python package missing: {module}")
+    if shutil.which(os.environ.get("DPA4SPIN_KOKKOS_MPIEXEC", "mpirun")) is None:
+        _unavailable("mpirun required for the two-rank /kk regression")
+    model = Path(os.environ.get("DPA4SPIN_KOKKOS_MODEL", _MODEL))
+    if not model.is_file():
+        _unavailable(f"Generate the compressed canonical fixture first: {model}")
+    with zipfile.ZipFile(model) as archive:
+        meta = json.loads(archive.read("model/extra/metadata.json"))
+    assert meta["lower_input_kind"] == "dpa4c_canonical"
+    assert meta["canonical_index_dtype"] == "auto"
+    assert meta["canonical_index_storage"] == "byte_words"
+    assert meta["graph_edge_dtype"] == "float32"
+    assert meta["spin_scheme"] == "native"
+    assert meta["type_map"] == ["Ni", "O"]
+    assert meta["use_spin"] == [True, False]
+    for key in ("dim_fparam", "dim_aparam", "dim_chg_spin"):
+        assert meta[key] == 0
+    assert not meta.get("has_comm_artifact", False)
+    return model
+
+
+def _run(tmp_path, data, model, *, ranks, kokkos):
+    output = tmp_path / f"{'kokkos' if kokkos else 'host'}-{ranks}.json"
+    runner = Path(__file__).with_name("run_mpi_dpa4spin_conventions.py")
+    command = [
+        os.environ.get("DPA4SPIN_KOKKOS_MPIEXEC", "mpirun"),
+        "-n",
+        str(ranks),
+        sys.executable,
+        str(runner),
+        str(data),
+        str(model),
+        str(output),
+        "--ranks",
+        str(ranks),
+    ]
+    if kokkos:
+        command.append("--kokkos")
+    # As in dpa4_spin_harness, bound the whole MPI process group: killing only
+    # mpirun can leave workers holding the pipes open after a timeout.
+    with subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        start_new_session=True,
+    ) as process:
+        try:
+            stdout, stderr = process.communicate(timeout=600)
+        except subprocess.TimeoutExpired:
+            try:
+                os.killpg(process.pid, signal.SIGKILL)
+            except ProcessLookupError:
+                # The process group already exited between timeout and cleanup.
+                pass
+            stdout, stderr = process.communicate()
+            pytest.fail(
+                f"MPI runner timed out (process group killed):\n{stdout}{stderr}"
+            )
+        assert process.returncode == 0, stdout + stderr
+    return json.loads(output.read_text())
+
+
+@pytest.fixture(scope="module", params=[0.0, 1e-8, 0.2], ids=["zero", "tiny", "finite"])
+def samples(kokkos_model, tmp_path_factory, request):
+    directory = tmp_path_factory.mktemp(f"dpa4spin-kk-{request.param}")
+    data = directory / "nio.lmp"
+    spin = _SPIN.copy()
+    spin[0] = request.param * np.array([2.0, -1.0, 2.0]) / 3.0
+    write_lmp_data_spin(_BOX, _COORD, spin, _TYPE, data)
+    return (
+        _run(directory, data, kokkos_model, ranks=1, kokkos=False),
+        _run(directory, data, kokkos_model, ranks=1, kokkos=True),
+        _run(directory, data, kokkos_model, ranks=2, kokkos=True),
+    )
+
+
+def _assert_snapshot(actual, expected, scale=1.0):
+    assert actual["ids"] == expected["ids"] == list(range(1, 7))
+    np.testing.assert_allclose(actual["spin"], expected["spin"], atol=0, rtol=0)
+    # The canonical kernels use float32; edge/reduction ordering differs between
+    # host, folded single-rank and extended multi-rank graphs. A snapshot holds
+    # the per-atom virial only when its run requested it; the global virial
+    # sums six atoms.
+    fields = [("force", 2e-7), ("magnetic_force", 5e-9), ("global_virial", 1e-6)]
+    if "virial" in expected:
+        fields.append(("virial", 2e-7))
+    for name, atol in fields:
+        np.testing.assert_allclose(
+            actual[name], scale * np.array(expected[name]), rtol=5e-4, atol=atol
+        )
+    if "virial" in actual:
+        # The pair fills the global and the per-atom virial from separate
+        # outputs of the model, whose totals agree. The first six centroid
+        # components (xx, yy, zz, xy, xz, yz) are the global ones.
+        np.testing.assert_allclose(
+            np.sum(actual["virial"], axis=0)[:6],
+            actual["global_virial"],
+            rtol=5e-4,
+            atol=1e-6,
+        )
+    assert actual["energy"] == pytest.approx(
+        scale * expected["energy"], rel=5e-6, abs=2e-7
+    )
+
+
+def test_kokkos_conventions_and_mpi(samples):
+    host, single, parallel = samples
+    reference = host["cartesian"]["snapshots"][0]
+    assert np.linalg.norm(np.array(reference["magnetic_force"])[0]) > 1e-7
+    for ranks, result in ((1, single), (2, parallel)):
+        cartesian = result["cartesian"]["snapshots"][0]
+        magnetic = np.array(cartesian["magnetic_force"])
+        # In particular, forcing the GPU kernel to precession at S=0 fails here.
+        assert np.linalg.norm(magnetic[0]) > 1e-7
+        np.testing.assert_array_equal(magnetic[3:], 0.0)
+        _assert_snapshot(cartesian, reference)
+        for selection in ("default", "precession", "cartesian"):
+            sample = result[selection]
+            assert sample["convention"] == int(selection == "cartesian")
+            assert sample["pair_style"] == "dpa4spin/kk"
+            assert sample["atom_style"] == "spin/kk"
+            assert sample["run_style_active"] == 1
+            assert sample["kokkos_active"] == sample["kokkos_ngpus"] == 1
+            assert sample["newton_pair"] == 0
+            assert len(sample["local_counts"]) == ranks
+            assert sum(sample["local_counts"]) == 6
+            assert min(sample["local_counts"]) > 0
+            if ranks == 2:
+                assert min(sample["ghost_counts"]) > 0
+            actual = sample["snapshots"][0]
+            _assert_snapshot(actual, host[selection]["snapshots"][0])
+            _assert_snapshot(actual, single[selection]["snapshots"][0])
+            if selection != "cartesian":
+                moments = np.array(actual["spin"])[:, 3]
+                finite = moments > 0
+                precession = np.array(actual["magnetic_force"])
+                # Normalize BEFORE comparing so the tiny-moment case cannot
+                # pass with zero force hidden by a fixed absolute tolerance.
+                np.testing.assert_allclose(
+                    precession[finite] * (HBAR_METAL / moments[finite, None]),
+                    magnetic[finite],
+                    rtol=5e-4,
+                    atol=5e-9,
+                )
+                np.testing.assert_array_equal(precession[~finite], 0.0)
+            for repeated in sample["snapshots"][1:]:
+                _assert_snapshot(repeated, actual)
+            virial_snapshots = sample["virial_snapshots"]
+            first = virial_snapshots[0]
+            _assert_snapshot(first, host[selection]["virial_snapshots"][0])
+            _assert_snapshot(first, single[selection]["virial_snapshots"][0])
+            # Requesting the per-atom virial changes no force or energy.
+            _assert_snapshot(first, actual)
+            _assert_snapshot(virial_snapshots[1], first)
+            _assert_snapshot(sample["scaled"], first, scale=2.75)
+
+
+def test_kokkos_overlay_accumulates(samples):
+    host, single, parallel = samples
+    for result in (single, parallel):
+        for selection in ("precession", "cartesian"):
+            overlay = result[f"overlay_{selection}"]
+            assert overlay["pair_style"] == "hybrid/overlay/kk"
+            assert overlay["run_style_active"] == 1
+            assert overlay["convention"] == int(selection == "cartesian")
+            for key in ("snapshots", "virial_snapshots"):
+                for snapshot in overlay[key]:
+                    _assert_snapshot(snapshot, host[selection][key][0], scale=2.0)
+                    _assert_snapshot(snapshot, result[selection][key][0], scale=2.0)
diff --git a/source/lmp/tests/test_dpa4spin_magnetic_force_convention.py b/source/lmp/tests/test_dpa4spin_magnetic_force_convention.py
new file mode 100644
index 0000000000..7855803905
--- /dev/null
+++ b/source/lmp/tests/test_dpa4spin_magnetic_force_convention.py
@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Exercise magnetic-force conventions through the registered LAMMPS adapter.
+
+Generate the CPU graph fixture with ``source/tests/infer/gen_dpa4c_spin_lmp.py``.
+PyTorch-enabled LAMMPS CI sets DPA4SPIN_LMP_REQUIRED=1 so missing fixtures,
+Python bindings, or plugin registration fail instead of hiding regressions.
+These tests exercise the host adapter; the GPU-only Kokkos adapter requires a
+separate compressed canonical archive and is not covered by this CPU fixture.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+import ctypes
+import importlib
+import os
+from contextlib import (
+    contextmanager,
+)
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+import pytest
+from write_lmp_data import (
+    write_lmp_data_spin,
+)
+
+HBAR_METAL = 6.5821191e-04
+_MODEL = (
+    Path(__file__).resolve().parents[2]
+    / "tests"
+    / "infer"
+    / "deeppot_dpa4c_spin_lmp.pt2"
+)
+_BOX = np.array([0.0, 6.0, 0.0, 6.0, 0.0, 6.0, 0.0, 0.0, 0.0])
+_COORD = np.array(
+    [
+        [1.0, 1.0, 1.0],
+        [3.2, 1.4, 1.1],
+        [1.3, 1.8, 1.0],
+        [0.4, 1.2, 1.6],
+        [3.6, 2.0, 1.3],
+        [3.4, 0.7, 1.7],
+    ]
+)
+_SPIN = np.array(
+    [
+        [0.0, 0.0, 0.0],
+        [-0.07, 0.09, 0.03],
+        [0.02, -0.06, 0.08],
+        [0.0, 0.0, 0.0],
+        [0.0, 0.0, 0.0],
+        [0.0, 0.0, 0.0],
+    ]
+)
+_TYPE = np.array([1, 1, 1, 2, 2, 2], dtype=np.int32)
+
+
+def _unavailable(reason: str) -> None:
+    if os.environ.get("DPA4SPIN_LMP_REQUIRED") == "1":
+        pytest.fail(reason)
+    pytest.skip(reason)
+
+
+@pytest.fixture(scope="module")
+def adapter():
+    if os.environ.get("ENABLE_PYTORCH", "1") != "1":
+        pytest.skip("PyTorch backend is disabled")
+    try:
+        module = importlib.import_module("lammps")
+    except ImportError as exc:
+        _unavailable(f"LAMMPS Python bindings unavailable: {exc}")
+        raise  # Keep the import failure path explicit for static analysis.
+    model = Path(os.environ.get("DPA4SPIN_LMP_MODEL", _MODEL))
+    if not model.is_file():
+        _unavailable(f"Generate the native-spin LAMMPS fixture first: {model}")
+    with module.lammps(cmdargs=["-log", "none", "-screen", "none"]) as lmp:
+        if not lmp.has_style("pair", "dpa4spin"):
+            _unavailable("LAMMPS plugin did not register pair style dpa4spin")
+    return module.lammps, model
+
+
+@pytest.fixture
+def data_file(tmp_path: Path) -> Path:
+    path = tmp_path / "nio.lmp"
+    write_lmp_data_spin(_BOX, _COORD, _SPIN, _TYPE, path)
+    return path
+
+
+@contextmanager
+def _system(adapter, data_file: Path, pair_args: str = "", units: str = "metal"):
+    lammps_cls, model = adapter
+    with lammps_cls(cmdargs=["-log", "none", "-screen", "none"]) as lmp:
+        for command in (
+            f"units {units}",
+            "atom_style spin",
+            "atom_modify map yes",
+            f"read_data {data_file.resolve()}",
+            "mass 1 58.6934",
+            "mass 2 15.9994",
+            f"pair_style dpa4spin {model.resolve()} {pair_args}",
+            "pair_coeff * * Ni O",
+            "neighbor 2.0 bin",
+            "compute fmprop all property/atom fmx fmy fmz",
+            "compute spprop all property/atom spx spy spz sp",
+            "run 0",
+        ):
+            lmp.command(command)
+        yield lmp
+
+
+def _pair_pointer(lmp, name: str, pointer_type):
+    # extract_pair() assumes double values, whereas the convention flag is int.
+    extract = lmp.lib.lammps_extract_pair
+    previous = extract.restype
+    extract.restype = pointer_type
+    try:
+        pointer = extract(lmp.lmp, name.encode())
+    finally:
+        extract.restype = previous
+    assert pointer, f"Pair::extract({name!r}) returned null"
+    return pointer
+
+
+def _convention(lmp) -> int:
+    assert lmp.extract_pair_dimension("fm_cartesian_force") == 0
+    return _pair_pointer(lmp, "fm_cartesian_force", ctypes.POINTER(ctypes.c_int))[0]
+
+
+def _gather(lmp, name: str, columns: int) -> np.ndarray:
+    return np.array(lmp.gather(name, 1, columns)).reshape(-1, columns)
+
+
+@pytest.mark.parametrize("magnitude", [0.0, 1e-8, 0.2], ids=["zero", "tiny", "finite"])
+def test_cartesian_matches_energy_gradient(adapter, data_file: Path, magnitude: float):
+    """Check every Cartesian component, including the derivative at S=0."""
+    spin = _SPIN.copy()
+    spin[0] = magnitude * np.array([2.0, -1.0, 2.0]) / 3.0
+    write_lmp_data_spin(_BOX, _COORD, spin, _TYPE, data_file)
+    with _system(adapter, data_file, "magnetic_force cartesian") as lmp:
+        assert _convention(lmp) == 1
+        analytic = _gather(lmp, "c_fmprop", 3)[0]
+        assert np.all(np.abs(analytic) > 1e-9), (
+            "Every tested spin component must be live",
+            analytic,
+        )
+        initial = _gather(lmp, "c_spprop", 4)[0]
+        moment = initial[:3] * initial[3]
+        eps = 1e-5
+        finite_difference = np.empty(3)
+        for component in range(3):
+            energies = []
+            for sign in (1, -1):
+                displaced = moment.copy()
+                displaced[component] += sign * eps
+                norm = np.linalg.norm(displaced)
+                direction = displaced / norm
+                # Both displaced norms are positive even when the centre is S=0.
+                # LAMMPS's set command rejects an exactly zero magnitude.
+                lmp.command(
+                    f"set atom 1 spin/atom {norm:.17g} "
+                    + " ".join(f"{value:.17g}" for value in direction)
+                )
+                lmp.command("run 0")
+                energies.append(lmp.get_thermo("pe"))
+            finite_difference[component] = -(energies[0] - energies[1]) / (2 * eps)
+        # Check components separately: a norm-only comparison could hide a
+        # sign error in the small transverse component of the zero-spin atom.
+        np.testing.assert_allclose(finite_difference, analytic, rtol=5e-3, atol=1e-11)
+
+
+def test_precession_default_and_explicit_match_cartesian(adapter, data_file: Path):
+    with _system(adapter, data_file, "magnetic_force cartesian") as lmp:
+        cartesian = _gather(lmp, "c_fmprop", 3)
+        moments = _gather(lmp, "c_spprop", 4)[:, 3]
+        force = _gather(lmp, "f", 3)
+        energy = lmp.get_thermo("pe")
+    assert np.linalg.norm(cartesian[0]) > 1e-8
+    np.testing.assert_allclose(cartesian[3:], 0.0, atol=1e-14)
+    for selection in ("", "magnetic_force precession"):
+        with _system(adapter, data_file, selection) as lmp:
+            assert _convention(lmp) == 0
+            actual = _gather(lmp, "c_fmprop", 3)
+            np.testing.assert_allclose(
+                actual,
+                cartesian * (moments / HBAR_METAL)[:, None],
+                rtol=1e-10,
+                atol=1e-14,
+            )
+            np.testing.assert_array_equal(actual[0], 0.0)
+            np.testing.assert_allclose(
+                _gather(lmp, "f", 3), force, rtol=1e-12, atol=1e-14
+            )
+            assert lmp.get_thermo("pe") == pytest.approx(energy, rel=1e-12)
+
+
+@pytest.mark.parametrize("selection", ["precession", "cartesian"])
+def test_pair_scale_and_repeated_compute(adapter, data_file: Path, selection: str):
+    """Exercise the actual extract/scale and force accumulation paths."""
+    with _system(adapter, data_file, f"magnetic_force {selection}") as lmp:
+        magnetic_force = _gather(lmp, "c_fmprop", 3)
+        force = _gather(lmp, "f", 3)
+        energy = lmp.get_thermo("pe")
+        scale = _pair_pointer(
+            lmp, "scale", ctypes.POINTER(ctypes.POINTER(ctypes.c_double))
+        )
+        scale[1][1] = 2.75
+        for _ in range(2):
+            lmp.command("run 0")
+            np.testing.assert_allclose(
+                _gather(lmp, "c_fmprop", 3),
+                2.75 * magnetic_force,
+                rtol=1e-10,
+                atol=1e-14,
+            )
+            np.testing.assert_allclose(
+                _gather(lmp, "f", 3), 2.75 * force, rtol=1e-10, atol=1e-14
+            )
+            assert lmp.get_thermo("pe") == pytest.approx(2.75 * energy, rel=1e-12)
+
+
+@pytest.mark.parametrize(
+    "pair_args, units, message",
+    [
+        ("magnetic_force", "metal", "takes one of precession or cartesian"),
+        ("magnetic_force unknown", "metal", "takes one of precession or cartesian"),
+        (
+            "magnetic_force cartesian magnetic_force precession",
+            "metal",
+            "magnetic_force may appear only once",
+        ),
+        ("magnetic_force cartesian", "real", "requires units metal"),
+    ],
+)
+def test_invalid_magnetic_force_settings(
+    adapter, data_file: Path, pair_args, units, message
+):
+    with (
+        pytest.raises(Exception, match=message),
+        _system(adapter, data_file, pair_args, units),
+    ):
+        pytest.fail("Invalid magnetic_force setting was accepted")
diff --git a/source/lmp/tests/test_lammps_dpa4c_long_range_kokkos.py b/source/lmp/tests/test_lammps_dpa4c_long_range_kokkos.py
new file mode 100644
index 0000000000..9765fdfe2b
--- /dev/null
+++ b/source/lmp/tests/test_lammps_dpa4c_long_range_kokkos.py
@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Test the Kokkos pair style with compressed DPA4C models that carry a long-range module.
+
+``pair_style deepmd/kk`` evaluates a compressed long-range model through the
+device-resident compact canonical lower: it gathers the positions of the
+model nodes relative to the origin of the box, passes the cell of a periodic
+system, and joins the exchange of the long range on every rank. Run
+source/install/test_lammps_kokkos.sh after test_cc_local.sh on a CUDA host:
+that job supplies the compressed fixtures of
+source/tests/infer/gen_dpa4c_long_range.py and a Kokkos LAMMPS library and
+plugin, and sets DPA4C_LONG_RANGE_KOKKOS_REQUIRED=1, so that no missing
+prerequisite skips the test silently.
+
+The six atoms of the fixtures, and their box, are translated away from the
+origin: the positions relative to the box, and with them every output of the
+references, including the per-atom virial, which depends on the origin of a
+periodic cell, stay those of the references.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+import importlib.util
+import os
+import shutil
+from pathlib import (
+    Path,
+)
+
+import constants
+import numpy as np
+import pytest
+from expected_ref import (
+    read_expected_ref,
+)
+from lammps_test_utils import (
+    run_mpi_pair_runner,
+)
+from write_lmp_data import (
+    write_lmp_data,
+)
+
+_INFER = Path(__file__).resolve().parents[2] / "tests" / "infer"
+_RUNNER = Path(__file__).with_name("run_mpi_pair_deepmd_dpa3_pt2.py")
+# The translation of the atoms and of the box.
+_ORIGIN = np.array([-4.0, 2.5, 7.0])
+_COORD = np.array(
+    [
+        [12.83, 2.56, 2.18],
+        [12.09, 2.87, 2.74],
+        [0.25, 3.32, 1.68],
+        [3.36, 3.00, 1.81],
+        [3.51, 2.51, 2.60],
+        [4.27, 3.22, 1.56],
+    ]
+)
+# atype [0, 1, 1, 0, 1, 1] (O, H, H, O, H, H) -> LAMMPS types [1, 2, 2, 1, 2, 2]
+_TYPE = np.array([1, 2, 2, 1, 2, 2])
+# The box of the non-periodic references, and the triclinic cell of the
+# periodic ones (gen_dpa4c_long_range.BOX): xlo xhi ylo yhi zlo zhi xy xz yz.
+_BOX = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0], dtype=float)
+_PERIODIC_BOX = np.array([0, 13, 0, 12.5, 0, 13.5, 0.6, -0.4, 0.8])
+# LAMMPS centroid/stress/atom reports the per-atom stress, the negative of the
+# per-atom virial, with the components xx, yy, zz, xy, xz, yz, yx, zx, zy; these
+# columns of it give the nine components of the reference virial.
+_TO_REFERENCE = [0, 6, 7, 3, 1, 8, 4, 5, 2]
+
+
+@pytest.fixture(scope="module")
+def kokkos_job() -> None:
+    # The ordinary LAMMPS suite also runs on CUDA, but uses a non-Kokkos wheel.
+    # Only the dedicated matching library/plugin job opts into this regression.
+    if os.environ.get("DPA4C_LONG_RANGE_KOKKOS_REQUIRED") != "1":
+        pytest.skip("Enable the dedicated CUDA Kokkos plugin regression explicitly")
+    for module in ("lammps", "mpi4py"):
+        if importlib.util.find_spec(module) is None:
+            pytest.fail(f"Required Python package missing: {module}")
+    if shutil.which("mpirun") is None:
+        pytest.fail("mpirun required for the Kokkos regression")
+
+
+def _translated(box: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """The box and the positions of the atoms, translated by ``_ORIGIN``."""
+    shifted = box.copy()
+    shifted[0:6] += np.repeat(_ORIGIN, 2)
+    return shifted, _COORD + _ORIGIN
+
+
+@pytest.mark.parametrize("nprocs", [1, 2])
+@pytest.mark.parametrize(
+    ("stem", "boundary", "section", "pair_extra"),
+    [
+        ("deeppot_dpa4c_long_range_compressed", "f f f", "default", ""),
+        (
+            "deeppot_dpa4c_long_range_compressed",
+            "f f f",
+            "explicit",
+            "charge_spin 1.0 2.0",
+        ),
+        ("deeppot_dpa4c_long_range_periodic_compressed", "p p p", "periodic", ""),
+        ("deeppot_dpa4c_long_range_periodic_compressed", "f f f", "cluster", ""),
+    ],
+)
+def test_pair_deepmd_kokkos_long_range(
+    kokkos_job: None,
+    tmp_path: Path,
+    stem: str,
+    boundary: str,
+    section: str,
+    pair_extra: str,
+    nprocs: int,
+) -> None:
+    model = _INFER / f"{stem}.pt2"
+    if not model.is_file():
+        pytest.fail(f"Generate the compressed long-range fixture first: {model}")
+    reference = read_expected_ref(_INFER / f"{stem}.expected")[section]
+    box, coord = _translated(_PERIODIC_BOX if "periodic" in stem else _BOX)
+    data = tmp_path / "six.lmp"
+    write_lmp_data(box, coord, _TYPE, data)
+    runner_args = ["--boundary", boundary, "--kokkos"]
+    if pair_extra:
+        runner_args += ["--pair-extra", pair_extra]
+    result = run_mpi_pair_runner(
+        _RUNNER,
+        data,
+        model,
+        nprocs=nprocs,
+        processors="1 1 1" if nprocs == 1 else "2 1 1",
+        runner_args=runner_args,
+        timeout=600,
+    )
+    assert result["pe"] == pytest.approx(
+        float(np.sum(reference["expected_e"])), rel=1e-6
+    )
+    np.testing.assert_allclose(
+        result["forces"], reference["expected_f"].reshape(6, 3), rtol=0, atol=1e-5
+    )
+    np.testing.assert_allclose(
+        result["virials"][:, _TO_REFERENCE] / constants.nktv2p,
+        -reference["expected_v"].reshape(6, 9),
+        rtol=0,
+        atol=1e-5,
+    )
diff --git a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py
new file mode 100644
index 0000000000..2048750aa4
--- /dev/null
+++ b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py
@@ -0,0 +1,439 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Test LAMMPS with a .pt2 DPA4C model that carries a long-range module.
+
+The latent-charge module sums over every pair of atoms of a complete
+non-periodic system, and over the images of the cell of a periodic one when
+it serves periodic systems, so ``pair_style deepmd`` runs a model whose module
+serves no periodic system with non-periodic boundaries only, and a model whose
+module serves periodic systems with boundaries periodic in every dimension or
+in none. A compressed model also runs on several ranks, whose atoms it gathers
+over MPI; an uncompressed one runs on a single rank. The models and the reference
+values are generated by source/tests/infer/gen_dpa4c_long_range.py, the
+compressed models for CUDA only: without a cell, at the default charge state
+and at ``charge_spin 1.0 2.0``, for a module that serves no periodic system,
+and in its cell and without one for a module that does.
+"""
+
+import importlib.util
+import os
+import shutil
+from pathlib import (
+    Path,
+)
+
+import constants
+import numpy as np
+import pytest
+from expected_ref import (
+    read_expected_ref,
+)
+from lammps import (
+    PyLammps,
+)
+from lammps_test_utils import (
+    make_atomic_lammps,
+    run_mpi_pair_runner,
+)
+from write_lmp_data import (
+    write_lmp_data,
+)
+
+pt2_file = (
+    Path(__file__).parent.parent.parent
+    / "tests"
+    / "infer"
+    / "deeppot_dpa4c_long_range.pt2"
+)
+ref_file = (
+    Path(__file__).parent.parent.parent
+    / "tests"
+    / "infer"
+    / "deeppot_dpa4c_long_range.expected"
+)
+compressed_pt2_file = pt2_file.with_name("deeppot_dpa4c_long_range_compressed.pt2")
+periodic_pt2_file = pt2_file.with_name("deeppot_dpa4c_long_range_periodic.pt2")
+periodic_compressed_pt2_file = pt2_file.with_name(
+    "deeppot_dpa4c_long_range_periodic_compressed.pt2"
+)
+# A model without a long-range module, for a model deviation beside one.
+short_range_file = pt2_file.with_name("deeppot_sea.pth")
+compressed_ref_file = ref_file.with_name("deeppot_dpa4c_long_range_compressed.expected")
+data_file = Path(__file__).parent / "data_dpa4c_long_range_pt2.lmp"
+periodic_data_file = Path(__file__).parent / "data_dpa4c_long_range_periodic_pt2.lmp"
+replicated_data_file = Path(__file__).parent / "data_dpa4c_long_range_grid_pt2.lmp"
+mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py"
+
+# Guarded with try/except because the generator only runs when PyTorch is
+# built; matrices that disable PyTorch still load this file at collection time.
+try:
+    _ref = read_expected_ref(ref_file)
+    expected_e_default = float(np.sum(_ref["default"]["expected_e"]))
+    expected_f_default = _ref["default"]["expected_f"].reshape(6, 3)
+    expected_e_explicit = float(np.sum(_ref["explicit"]["expected_e"]))
+    expected_f_explicit = _ref["explicit"]["expected_f"].reshape(6, 3)
+except FileNotFoundError:
+    expected_e_default = expected_f_default = None
+    expected_e_explicit = expected_f_explicit = None
+
+
+def _read_ref(path: Path) -> dict | None:
+    try:
+        return read_expected_ref(path)
+    except FileNotFoundError:
+        return None
+
+
+compressed_ref = _read_ref(compressed_ref_file)
+periodic_ref = _read_ref(
+    ref_file.with_name("deeppot_dpa4c_long_range_periodic.expected")
+)
+periodic_compressed_ref = _read_ref(
+    ref_file.with_name("deeppot_dpa4c_long_range_periodic_compressed.expected")
+)
+
+box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0])
+# The triclinic cell of the periodic references: xlo xhi ylo yhi zlo zhi xy xz
+# yz of the lattice vectors (13, 0, 0), (0.6, 12.5, 0) and (-0.4, 0.8, 13.5).
+periodic_box = np.array([0, 13, 0, 12.5, 0, 13.5, 0.6, -0.4, 0.8])
+coord = np.array(
+    [
+        [12.83, 2.56, 2.18],
+        [12.09, 2.87, 2.74],
+        [0.25, 3.32, 1.68],
+        [3.36, 3.00, 1.81],
+        [3.51, 2.51, 2.60],
+        [4.27, 3.22, 1.56],
+    ]
+)
+# atype [0, 1, 1, 0, 1, 1] (O, H, H, O, H, H) -> LAMMPS types [1, 2, 2, 1, 2, 2]
+type_OH = np.array([1, 2, 2, 1, 2, 2])
+# LAMMPS centroid/stress/atom reports the per-atom stress, the negative of the
+# per-atom virial, with the components xx, yy, zz, xy, xz, yz, yx, zx, zy; these
+# columns of it give the nine components of the reference virial.
+expected_v_to_lammps = [0, 6, 7, 3, 1, 8, 4, 5, 2]
+
+
+def setup_module() -> None:
+    if os.environ.get("ENABLE_PYTORCH", "1") != "1":
+        pytest.skip("Skip test because PyTorch support is not enabled.")
+    write_lmp_data(box, coord, type_OH, data_file)
+    write_lmp_data(periodic_box, coord, type_OH, periodic_data_file)
+
+
+def teardown_module() -> None:
+    for path in (data_file, periodic_data_file, replicated_data_file):
+        if path.exists():
+            os.remove(path)
+
+
+def _lammps(boundary: str, data: Path = data_file) -> PyLammps:
+    if not pt2_file.is_file():
+        pytest.skip(f"Skip because {pt2_file} was not generated.")
+    return make_atomic_lammps(data, boundary=boundary, atom_map="yes")
+
+
+@pytest.fixture
+def lammps():
+    lmp = _lammps("f f f")
+    yield lmp
+    lmp.close()
+
+
+@pytest.fixture
+def lammps_periodic():
+    lmp = _lammps("p p p")
+    yield lmp
+    lmp.close()
+
+
+def test_pair_deepmd(lammps) -> None:
+    lammps.pair_style(f"deepmd {pt2_file.resolve()}")
+    lammps.pair_coeff("* *")
+    lammps.run(0)
+    assert lammps.eval("pe") == pytest.approx(expected_e_default)
+    for ii in range(6):
+        assert lammps.atoms[ii].force == pytest.approx(
+            expected_f_default[lammps.atoms[ii].id - 1]
+        )
+    lammps.run(1)
+
+
+def test_pair_deepmd_charge_spin(lammps) -> None:
+    lammps.pair_style(f"deepmd {pt2_file.resolve()} charge_spin 1.0 2.0")
+    lammps.pair_coeff("* *")
+    lammps.run(0)
+    assert lammps.eval("pe") == pytest.approx(expected_e_explicit)
+    for ii in range(6):
+        assert lammps.atoms[ii].force == pytest.approx(
+            expected_f_explicit[lammps.atoms[ii].id - 1]
+        )
+
+
+def test_periodic_boundary_is_rejected(lammps_periodic) -> None:
+    lammps_periodic.pair_style(f"deepmd {pt2_file.resolve()}")
+    lammps_periodic.pair_coeff("* *")
+    with pytest.raises(Exception, match="long-range module"):
+        lammps_periodic.run(0)
+
+
+@pytest.mark.parametrize("long_range_first", [True, False])
+def test_model_deviation_periodic_boundary_is_rejected(
+    lammps_periodic, long_range_first: bool, tmp_path: Path
+) -> None:
+    """Any model of a model deviation with a long-range module confines it."""
+    if not short_range_file.is_file():
+        pytest.skip(f"Skip because {short_range_file} was not generated.")
+    models = [pt2_file, short_range_file]
+    if not long_range_first:
+        models.reverse()
+    lammps_periodic.pair_style(
+        "deepmd "
+        + " ".join(str(model.resolve()) for model in models)
+        + f" out_freq 1 out_file {tmp_path / 'model_devi.out'}"
+    )
+    lammps_periodic.pair_coeff("* *")
+    with pytest.raises(Exception, match="long-range module"):
+        lammps_periodic.run(0)
+
+
+@pytest.mark.parametrize("long_range_first", [True, False])
+def test_model_deviation_without_cell_needs_long_range_throughout(
+    lammps, long_range_first: bool, tmp_path: Path
+) -> None:
+    """A non-periodic model deviation passes no cell, which every model must serve."""
+    if not short_range_file.is_file():
+        pytest.skip(f"Skip because {short_range_file} was not generated.")
+    models = [pt2_file, short_range_file]
+    if not long_range_first:
+        models.reverse()
+    lammps.pair_style(
+        "deepmd "
+        + " ".join(str(model.resolve()) for model in models)
+        + f" out_freq 1 out_file {tmp_path / 'model_devi.out'}"
+    )
+    lammps.pair_coeff("* *")
+    with pytest.raises(Exception, match="every model carries one"):
+        lammps.run(0)
+
+
+def test_model_deviation_of_long_range_models(lammps, tmp_path: Path) -> None:
+    """Models that all carry a long-range module serve a non-periodic deviation."""
+    md_file = tmp_path / "model_devi.out"
+    lammps.pair_style(
+        f"deepmd {pt2_file.resolve()} {pt2_file.resolve()} "
+        f"out_freq 1 out_file {md_file}"
+    )
+    lammps.pair_coeff("* *")
+    lammps.run(0)
+    assert lammps.eval("pe") == pytest.approx(expected_e_default)
+    for ii in range(6):
+        assert lammps.atoms[ii].force == pytest.approx(
+            expected_f_default[lammps.atoms[ii].id - 1]
+        )
+    md = np.loadtxt(md_file)
+    assert md[1:] == pytest.approx(np.zeros(md.shape[0] - 1), abs=1e-10)
+
+
+def _check_reference(lmp: PyLammps, ref: dict, section: str) -> None:
+    """Compare the energy and the forces of a run with a reference section."""
+    lmp.run(0)
+    assert lmp.eval("pe") == pytest.approx(float(np.sum(ref[section]["expected_e"])))
+    expected_f = ref[section]["expected_f"].reshape(6, 3)
+    for ii in range(6):
+        assert lmp.atoms[ii].force == pytest.approx(
+            expected_f[lmp.atoms[ii].id - 1], abs=1e-8
+        )
+
+
+@pytest.mark.parametrize(
+    ("boundary", "section"), [("p p p", "periodic"), ("f f f", "cluster")]
+)
+def test_pair_deepmd_periodic_module(boundary: str, section: str) -> None:
+    """A module that serves periodic systems sums over the images of the cell."""
+    if periodic_ref is None or not periodic_pt2_file.is_file():
+        pytest.skip(f"Skip because {periodic_pt2_file} was not generated.")
+    lmp = _lammps(boundary, periodic_data_file)
+    lmp.pair_style(f"deepmd {periodic_pt2_file.resolve()}")
+    lmp.pair_coeff("* *")
+    _check_reference(lmp, periodic_ref, section)
+    lmp.run(1)
+    lmp.close()
+
+
+@pytest.mark.parametrize("model", [pt2_file, periodic_pt2_file])
+def test_partially_periodic_boundary_is_rejected(model: Path) -> None:
+    if not model.is_file():
+        pytest.skip(f"Skip because {model} was not generated.")
+    lmp = _lammps("p p f", periodic_data_file)
+    lmp.pair_style(f"deepmd {model.resolve()}")
+    lmp.pair_coeff("* *")
+    with pytest.raises(Exception, match="every dimension"):
+        lmp.run(0)
+    lmp.close()
+
+
+mpi_available = pytest.mark.skipif(
+    shutil.which("mpirun") is None or importlib.util.find_spec("mpi4py") is None,
+    reason="MPI or mpi4py is not installed",
+)
+
+
+@mpi_available
+@pytest.mark.parametrize(
+    ("nprocs", "processors", "section", "pair_extra"),
+    [
+        # Both ranks own atoms.
+        (2, "2 1 1", "default", ""),
+        # Every atom lies in the first slab along z: two ranks hold ghost atoms
+        # alone and the last none at all, and all of them join the exchange.
+        (4, "1 1 4", "explicit", "charge_spin 1.0 2.0"),
+    ],
+)
+def test_pair_deepmd_mpi_compressed(
+    nprocs: int, processors: str, section: str, pair_extra: str
+) -> None:
+    if compressed_ref is None or not compressed_pt2_file.is_file():
+        pytest.skip(f"Skip because {compressed_pt2_file} was not generated.")
+    runner_args = ["--boundary", "f f f"]
+    if pair_extra:
+        runner_args += ["--pair-extra", pair_extra]
+    result = run_mpi_pair_runner(
+        mpi_runner,
+        data_file,
+        compressed_pt2_file,
+        nprocs=nprocs,
+        processors=processors,
+        runner_args=runner_args,
+    )
+    expected_e = float(np.sum(compressed_ref[section]["expected_e"]))
+    expected_f = compressed_ref[section]["expected_f"].reshape(6, 3)
+    expected_v = -compressed_ref[section]["expected_v"].reshape(6, 9)
+    assert result["pe"] == pytest.approx(expected_e, rel=1e-6)
+    np.testing.assert_allclose(result["forces"], expected_f, rtol=0, atol=1e-5)
+    np.testing.assert_allclose(
+        result["virials"][:, expected_v_to_lammps] / constants.nktv2p,
+        expected_v,
+        rtol=0,
+        atol=1e-5,
+    )
+
+
+@mpi_available
+@pytest.mark.parametrize(
+    ("nprocs", "processors", "boundary", "section"),
+    [
+        (2, "2 1 1", "p p p", "periodic"),
+        # Every atom lies in the lower half along z; the ranks above hold
+        # ghost atoms alone.
+        (4, "1 1 4", "p p p", "periodic"),
+        (2, "1 2 1", "f f f", "cluster"),
+    ],
+)
+def test_pair_deepmd_mpi_periodic_compressed(
+    nprocs: int, processors: str, boundary: str, section: str
+) -> None:
+    if periodic_compressed_ref is None or not periodic_compressed_pt2_file.is_file():
+        pytest.skip(f"Skip because {periodic_compressed_pt2_file} was not generated.")
+    if not periodic_data_file.is_file():
+        write_lmp_data(periodic_box, coord, type_OH, periodic_data_file)
+    result = run_mpi_pair_runner(
+        mpi_runner,
+        periodic_data_file,
+        periodic_compressed_pt2_file,
+        nprocs=nprocs,
+        processors=processors,
+        runner_args=["--boundary", boundary],
+    )
+    ref = periodic_compressed_ref[section]
+    expected_v = -ref["expected_v"].reshape(6, 9)
+    assert result["pe"] == pytest.approx(float(np.sum(ref["expected_e"])), rel=1e-6)
+    np.testing.assert_allclose(
+        result["forces"], ref["expected_f"].reshape(6, 3), rtol=0, atol=1e-5
+    )
+    np.testing.assert_allclose(
+        result["virials"][:, expected_v_to_lammps] / constants.nktv2p,
+        expected_v,
+        rtol=0,
+        atol=1e-5,
+    )
+
+
+@pytest.fixture(scope="module")
+def replicated_reference() -> dict:
+    """A grid of 168 copies of the six atoms, and its single-rank result.
+
+    The copies lie farther apart than the short-range cutoff, so only the
+    long-range module couples them; the 1008 atoms fill several blocks and
+    chunks of the long-range operator.
+    """
+    if compressed_ref is None or not compressed_pt2_file.is_file():
+        pytest.skip(f"Skip because {compressed_pt2_file} was not generated.")
+    shifts = np.stack(
+        np.meshgrid(
+            18.0 * np.arange(4), 6.0 * np.arange(7), 6.0 * np.arange(6), indexing="ij"
+        ),
+        axis=-1,
+    ).reshape(-1, 1, 3)
+    grid_coord = (coord[None] + shifts).reshape(-1, 3)
+    grid_type = np.tile(type_OH, shifts.shape[0])
+    # The box reaches beyond the atoms along z, so that the last of four slabs
+    # holds no atom, owned or ghost.
+    grid_box = np.array([0, 72, 0, 42, 0, 60, 0, 0, 0])
+    write_lmp_data(grid_box, grid_coord, grid_type, replicated_data_file)
+    return run_mpi_pair_runner(
+        mpi_runner,
+        replicated_data_file,
+        compressed_pt2_file,
+        nprocs=1,
+        runner_args=["--boundary", "f f f"],
+    )
+
+
+@mpi_available
+@pytest.mark.parametrize(
+    ("nprocs", "processors"),
+    [
+        # Uneven slabs that cut through copies of the six atoms.
+        (3, "3 1 1"),
+        # The last slab holds no atom at all.
+        (4, "1 1 4"),
+    ],
+)
+def test_pair_deepmd_mpi_compressed_matches_single_rank(
+    replicated_reference: dict, nprocs: int, processors: str
+) -> None:
+    result = run_mpi_pair_runner(
+        mpi_runner,
+        replicated_data_file,
+        compressed_pt2_file,
+        nprocs=nprocs,
+        processors=processors,
+        runner_args=["--boundary", "f f f"],
+    )
+    assert result["pe"] == pytest.approx(replicated_reference["pe"], rel=1e-6)
+    np.testing.assert_allclose(
+        result["forces"], replicated_reference["forces"], rtol=0, atol=1e-5
+    )
+    np.testing.assert_allclose(
+        result["virials"] / constants.nktv2p,
+        replicated_reference["virials"] / constants.nktv2p,
+        rtol=0,
+        atol=1e-5,
+    )
+
+
+@mpi_available
+def test_pair_deepmd_mpi_uncompressed_is_rejected() -> None:
+    if not pt2_file.is_file():
+        pytest.skip(f"Skip because {pt2_file} was not generated.")
+    result = run_mpi_pair_runner(
+        mpi_runner,
+        data_file,
+        pt2_file,
+        nprocs=2,
+        processors="2 1 1",
+        runner_args=["--boundary", "f f f"],
+        capture=True,
+    )
+    assert result["returncode"] != 0
+    assert "compact canonical lower" in result["stdout"] + result["stderr"]
diff --git a/source/lmp/tests/test_lammps_plugin.py b/source/lmp/tests/test_lammps_plugin.py
new file mode 100644
index 0000000000..d7073fe1d4
--- /dev/null
+++ b/source/lmp/tests/test_lammps_plugin.py
@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Check style registration without requiring a model or inference backend.
+
+CI loads the runtime plugin through LAMMPS_PLUGIN_PATH. A standalone plugin
+can instead be loaded explicitly with DEEPMD_TEST_PLUGIN.
+"""
+
+import os
+
+from lammps import (
+    lammps,
+)
+
+
+def test_dpa4spin_host_style_registered() -> None:
+    """The unsuffixed host style must also be available without Kokkos."""
+    with lammps(cmdargs=["-log", "none", "-screen", "none"]) as lmp:
+        if plugin := os.environ.get("DEEPMD_TEST_PLUGIN"):
+            lmp.command(f'plugin load "{plugin}"')
+        assert lmp.has_style("pair", "dpa4spin"), (
+            "LAMMPS plugin did not register pair style dpa4spin"
+        )
diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt
index 5702e6c18c..4c7e4e29d1 100644
--- a/source/op/pt/CMakeLists.txt
+++ b/source/op/pt/CMakeLists.txt
@@ -23,7 +23,8 @@ list(
   cpu/neighbor_search_cpu.cc
   dpa4c/ops.cc
   dpa4c/graph_compress_cpu.cc
-  ${DPA4C_CPU_ISA_SRC})
+  ${DPA4C_CPU_ISA_SRC}
+  long_range/ops.cc)
 # The architecture level is pinned per unit rather than added to it: the project
 # may be configured with ENABLE_NATIVE_OPTIMIZATION, and a trailing -march wins,
 # so this is what keeps the fallback units free of instructions the dispatcher
@@ -51,6 +52,15 @@ endif()
 # only within its own subtree, which does not cover this target.
 if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA)
   find_package(CUDAToolkit REQUIRED)
+  # The fused fitting network launches thread-block clusters through
+  # cudaLaunchKernelEx and the cluster occupancy queries, which CUDA 11.8
+  # introduced.
+  if(CUDAToolkit_VERSION VERSION_LESS "11.8")
+    message(
+      FATAL_ERROR
+        "The CUDA operators of the PyTorch backend require CUDA 11.8 or later, "
+        "found CUDA ${CUDAToolkit_VERSION}")
+  endif()
   option(DEEPMD_ENABLE_DPA1_HIGH_LMAX
          "Instantiate experimental DPA1 CUDA kernels for lmax greater than one"
          OFF)
@@ -162,10 +172,62 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA)
   unset(DPA4_TRAIN_TYPE_bf16)
   unset(DPA4_TRAIN_TYPE_f64)
   unset(DPA4_TRAIN_DTYPES)
-  set(DPA4C_GRAPH_COMPRESS_KERNEL_SRC
-      dpa4c/graph_compress.cu dpa4c/graph_compress_c8.cu
-      dpa4c/graph_compress_c16.cu dpa4c/graph_compress_c32.cu
-      dpa4c/graph_compress_c64.cu dpa4c/graph_compress_c128.cu)
+  # The compressed DPA4C descriptor kernels form a compile-time grid over the
+  # scalar width and the angular degree, both directions of a degree sharing
+  # most of their code. They follow the generated-shard layout of the DPA4
+  # training kernels, so that the widest profiles, whose degrees take the
+  # longest to compile, do not serialize behind one nvcc invocation: 64 and 128
+  # channels take one source per degree. A narrower width keeps its degrees in
+  # one source, whose kernels share most of their code, so that its fatbinary,
+  # compressed as one, takes about a third of the size of three. The widths are
+  # those of DPA4C_FOR_EACH_CHANNEL and the degrees those that
+  # dpa4c/graph_compress.cu dispatches.
+  set(DPA4C_GRAPH_COMPRESS_KERNEL_SRC dpa4c/graph_compress.cu)
+  set(DPA4C_GENERATED_DIR "${CMAKE_CURRENT_BINARY_DIR}/dpa4c")
+  set(DPA4C_KERNEL_HEADER
+      "${CMAKE_CURRENT_SOURCE_DIR}/dpa4c/graph_compress_kernel.cuh")
+  file(MAKE_DIRECTORY "${DPA4C_GENERATED_DIR}")
+  foreach(DPA4C_SHARD_WIDTH IN ITEMS 8 16 32 64 128)
+    if(DPA4C_SHARD_WIDTH GREATER_EQUAL 64)
+      set(DPA4C_SHARD_GROUPS 2 3 4)
+    else()
+      set(DPA4C_SHARD_GROUPS 234)
+    endif()
+    foreach(DPA4C_SHARD_GROUP IN LISTS DPA4C_SHARD_GROUPS)
+      set(DPA4C_SHARD_INSTANTIATIONS)
+      string(LENGTH "${DPA4C_SHARD_GROUP}" DPA4C_SHARD_DEGREES)
+      math(EXPR DPA4C_SHARD_LAST "${DPA4C_SHARD_DEGREES} - 1")
+      foreach(DPA4C_SHARD_INDEX RANGE ${DPA4C_SHARD_LAST})
+        string(SUBSTRING "${DPA4C_SHARD_GROUP}" ${DPA4C_SHARD_INDEX} 1
+                         DPA4C_SHARD_LMAX)
+        foreach(DPA4C_SHARD_DIRECTION IN ITEMS forward backward)
+          string(
+            APPEND
+            DPA4C_SHARD_INSTANTIATIONS
+            "template void launch_${DPA4C_SHARD_DIRECTION}<"
+            "${DPA4C_SHARD_WIDTH}, ${DPA4C_SHARD_LMAX}>(\n"
+            "    const Arguments& arguments, cudaStream_t stream);\n")
+        endforeach()
+      endforeach()
+      set(DPA4C_SHARD
+          "${DPA4C_GENERATED_DIR}/c${DPA4C_SHARD_WIDTH}_l${DPA4C_SHARD_GROUP}.cu"
+      )
+      configure_file(dpa4c/shard.cu.in "${DPA4C_SHARD}" @ONLY)
+      list(APPEND DPA4C_GRAPH_COMPRESS_KERNEL_SRC "${DPA4C_SHARD}")
+    endforeach()
+  endforeach()
+  unset(DPA4C_SHARD)
+  unset(DPA4C_SHARD_DIRECTION)
+  unset(DPA4C_SHARD_LMAX)
+  unset(DPA4C_SHARD_INDEX)
+  unset(DPA4C_SHARD_LAST)
+  unset(DPA4C_SHARD_DEGREES)
+  unset(DPA4C_SHARD_INSTANTIATIONS)
+  unset(DPA4C_SHARD_GROUP)
+  unset(DPA4C_SHARD_GROUPS)
+  unset(DPA4C_SHARD_WIDTH)
+  unset(DPA4C_KERNEL_HEADER)
+  unset(DPA4C_GENERATED_DIR)
   set(DPA1_GRAPH_COMPRESS_KERNEL_SRC
       dpa1_graph_compress_c8.cu dpa1_graph_compress_c16.cu
       dpa1_graph_compress_c32.cu dpa1_graph_compress_c64.cu
@@ -179,6 +241,10 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA)
     ${DPA4C_GRAPH_COMPRESS_KERNEL_SRC}
     graph_fitting.cu
     edge_force_virial.cu
+    long_range/direct_sum.cu
+    long_range/domain_exchange.cc
+    long_range/latent_charge_sum.cu
+    long_range/mesh_sum.cu
     dpa1_graph_energy_force.cu
     dpa4/so2_conv.cu
     ${DPA4_SO2_CONV_KERNEL_SRC}
@@ -213,7 +279,7 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA)
     set_property(TARGET deepmd_op_pt PROPERTY CUDA_ARCHITECTURES
                                               ${DEEPMD_PT_CUDA_ARCHITECTURES})
   endif()
-  target_link_libraries(deepmd_op_pt PRIVATE CUDA::cublas CUDA::cublasLt)
+  target_link_libraries(deepmd_op_pt PRIVATE CUDA::cublasLt CUDA::cufft)
   # libtorch headers require at least C++17, and PyTorch 2.14+ requires C++20.
   # Preserve an explicitly higher CUDA standard requested by the user.
   set(DEEPMD_PT_CUDA_STANDARD 17)
@@ -239,12 +305,25 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA)
   set_source_files_properties(
     ${DPA4_SO2_CONV_KERNEL_SRC} dpa1_graph_descriptor.cu
     PROPERTIES COMPILE_OPTIONS "--threads=2")
-  # The compressed DPA1 and DPA4C kernels are instantiated one translation unit
-  # per channel width so their angular-degree and topology specializations
-  # compile in parallel.
+  # The compressed DPA1 kernels, instantiated one translation unit per channel
+  # width, and the compressed DPA4C sources -- the generated shards and the
+  # bindings, whose own kernels use double arithmetic and fmaxf alone -- are
+  # compiled with fast math.
   set_source_files_properties(
     ${DPA1_GRAPH_COMPRESS_KERNEL_SRC} ${DPA4C_GRAPH_COMPRESS_KERNEL_SRC}
     PROPERTIES COMPILE_OPTIONS "--use_fast_math")
+  # The DPA4C dispatch names one instantiation per width and degree, which only
+  # the generated shards define. The module is loaded at run time, where a
+  # missing definition would make the loading fail and take every operator of
+  # the library with it, so the link resolves every symbol instead.
+  if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    target_link_options(deepmd_op_pt PRIVATE "LINKER:--no-undefined")
+  endif()
+  # The pair loop of the long-range direct sum calls exp2f and rsqrtf, which
+  # fast math, without denormal handling, compiles to a single special-function
+  # instruction each.
+  set_source_files_properties(long_range/direct_sum.cu
+                              PROPERTIES COMPILE_OPTIONS "--use_fast_math")
 endif()
 if(${OP_CXX_ABI_PT} EQUAL ${OP_CXX_ABI})
   target_link_libraries(deepmd_op_pt PRIVATE ${LIB_DEEPMD})
diff --git a/source/op/pt/canonical_index.h b/source/op/pt/canonical_index.h
new file mode 100644
index 0000000000..2799b2a466
--- /dev/null
+++ b/source/op/pt/canonical_index.h
@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+#pragma once
+
+#include 
+
+#include 
+
+namespace deepmd {
+
+/**
+ * Restore integer words from an exported canonical index buffer without a copy.
+ *
+ * AOTInductor fixes input dtypes, so the deployment boundary carries integer
+ * words as bytes with a final axis of length four or eight. Its symbolic extent
+ * selects the integer width at runtime; the descriptor kernels receive ordinary
+ * typed tensors. Eager calls may supply those typed tensors directly.
+ */
+inline torch::Tensor canonical_index_view(const torch::Tensor& buffer) {
+  if (buffer.scalar_type() != torch::kUInt8) {
+    return buffer;
+  }
+  TORCH_CHECK(
+      buffer.dim() >= 2 && buffer.is_contiguous(),
+      "canonical index bytes must be contiguous with a final word axis");
+  const auto bytes = buffer.size(-1);
+  TORCH_CHECK(bytes == 4 || bytes == 8,
+              "canonical index word width must be 4 or 8 bytes, got ", bytes);
+  TORCH_CHECK(reinterpret_cast(buffer.data_ptr()) % bytes == 0,
+              "canonical index bytes must be aligned to their word width");
+  const auto dtype = bytes == 8 ? torch::kInt64 : torch::kUInt32;
+  return buffer.view(dtype).squeeze(-1);
+}
+
+}  // namespace deepmd
diff --git a/source/op/pt/cpu/edge_force_virial_cpu.cc b/source/op/pt/cpu/edge_force_virial_cpu.cc
index e271269165..97352c27e0 100644
--- a/source/op/pt/cpu/edge_force_virial_cpu.cc
+++ b/source/op/pt/cpu/edge_force_virial_cpu.cc
@@ -355,15 +355,24 @@ canonical_edge_force_virial(torch::Tensor edge_gradient,
                             bool want_atom_virial) {
   TORCH_CHECK(edge_gradient.device().is_cpu(),
               "canonical_edge_force_virial: the CPU kernel needs CPU tensors");
+  const int64_t node_count = node_capacity.expect_int();
+  TORCH_CHECK(destination_row_ptr.numel() == node_count + 1 &&
+                  source_row_ptr.numel() == node_count + 1,
+              "canonical_edge_force_virial: row pointers must have N + 1 "
+              "entries");
+  TORCH_CHECK(
+      source_order.dim() == 1 && source_order.numel() == edge_gradient.size(0),
+      "canonical_edge_force_virial: source_order must hold one "
+      "position per edge slot, expected ",
+      edge_gradient.size(0), ", got ", source_order.numel());
   const auto empty_index = torch::empty({0}, source_order.options());
   const auto empty_mask =
       torch::empty({0}, edge_gradient.options().dtype(torch::kBool));
-  return assemble_entry(node_capacity.expect_int(), edge_gradient.contiguous(),
-                        edge_vec.contiguous(), empty_mask, empty_index,
-                        destination_row_ptr.contiguous(),
-                        source_order.contiguous(), source_row_ptr.contiguous(),
-                        n_node_per_frame, edge_spin_gradient.contiguous(),
-                        want_atom_virial);
+  return assemble_entry(
+      node_count, edge_gradient.contiguous(), edge_vec.contiguous(), empty_mask,
+      empty_index, destination_row_ptr.contiguous(), source_order.contiguous(),
+      source_row_ptr.contiguous(), n_node_per_frame,
+      edge_spin_gradient.contiguous(), want_atom_virial);
 }
 
 torch::Tensor frame_scalar_sum(torch::Tensor node_scalar,
diff --git a/source/op/pt/dpa1_graph_compress.cu b/source/op/pt/dpa1_graph_compress.cu
index 7dc14cc83e..32e359bc7f 100644
--- a/source/op/pt/dpa1_graph_compress.cu
+++ b/source/op/pt/dpa1_graph_compress.cu
@@ -16,6 +16,7 @@
 #include 
 #include 
 
+#include "canonical_index.h"
 #include "dpa1_graph_compress_launch.h"
 
 #ifndef DEEPMD_ENABLE_DPA1_HIGH_LMAX
@@ -401,6 +402,7 @@ std::tuple dpa1_canonical_compress(
     double protection,
     double neighbors,
     int64_t basis_dim) {
+  source = deepmd::canonical_index_view(source);
   TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0),
               "dpa1_canonical_compress: source and edge_vec storage must "
               "share the edge axis");
@@ -442,6 +444,7 @@ torch::Tensor dpa1_canonical_compress_backward(
     double rcut_smooth,
     double protection,
     double neighbors) {
+  source = deepmd::canonical_index_view(source);
   TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0),
               "dpa1_canonical_compress_backward: source and edge_vec storage "
               "must share the edge axis");
diff --git a/source/op/pt/dpa4/so2_conv.cu b/source/op/pt/dpa4/so2_conv.cu
index 01af267c57..4c67cfd807 100644
--- a/source/op/pt/dpa4/so2_conv.cu
+++ b/source/op/pt/dpa4/so2_conv.cu
@@ -21,6 +21,7 @@
 #include 
 #include 
 
+#include "../fp32_ridge.h"
 #include "so2_conv_launch.h"
 
 namespace {
@@ -514,34 +515,6 @@ torch::Tensor dpa4_wigner_runs(torch::Tensor quat,
                     mono_exp.contiguous(), at::cuda::getCurrentCUDAStream());
 }
 
-/// Ridge point of the current device: peak fp32 FMA throughput over DRAM
-/// bandwidth, in FLOP per byte.
-///
-/// The fused convolution trades memory traffic for float32 SIMT arithmetic,
-/// so the arithmetic budget at which it stops paying scales with this ratio.
-/// The Python routing gate normalizes its measured threshold by the ridge of
-/// the card it was calibrated on, which makes the decision follow the actual
-/// part rather than an architecture name. Attribute queries cover every
-/// supported toolkit; the FMA-per-SM width is 128 lanes on every consumer and
-/// data-center part since Ampere except the A100 die (64), which the
-/// compute-capability pair identifies.
-double dpa4_fp32_ridge() {
-  int dev = 0;
-  cudaGetDevice(&dev);
-  int sm_count = 0, clock_khz = 0, mem_clock_khz = 0, bus_width = 0;
-  int major = 0, minor = 0;
-  cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
-  cudaDeviceGetAttribute(&clock_khz, cudaDevAttrClockRate, dev);
-  cudaDeviceGetAttribute(&mem_clock_khz, cudaDevAttrMemoryClockRate, dev);
-  cudaDeviceGetAttribute(&bus_width, cudaDevAttrGlobalMemoryBusWidth, dev);
-  cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev);
-  cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev);
-  const int lanes = (major < 8 || (major == 8 && minor == 0)) ? 64 : 128;
-  const double flops = 2.0 * lanes * sm_count * (clock_khz * 1e3);
-  const double bytes = (mem_clock_khz * 1e3) * (bus_width / 8.0) * 2.0;
-  return bytes > 0.0 ? flops / bytes : 0.0;
-}
-
 torch::Tensor dpa4_wigner_runs_backward(torch::Tensor grad_runs,
                                         torch::Tensor quat,
                                         torch::Tensor dmono_coeff,
@@ -586,7 +559,7 @@ TORCH_LIBRARY_FRAGMENT(deepmd, m) {
       "dpa4_wigner_runs_backward(Tensor grad_runs, Tensor quat, "
       "Tensor dmono_coeff, Tensor dmono_exp) -> Tensor");
   m.impl("dpa4_wigner_runs_backward", torch::kCUDA, &dpa4_wigner_runs_backward);
-  m.def("dpa4_fp32_ridge() -> float", &dpa4_fp32_ridge);
+  m.def("dpa4_fp32_ridge() -> float", &fp32_ridge);
 }
 
 namespace dpa4 {
diff --git a/source/op/pt/dpa4c/graph_compress.cu b/source/op/pt/dpa4c/graph_compress.cu
index 9a3e05a33a..48a8aaf062 100644
--- a/source/op/pt/dpa4c/graph_compress.cu
+++ b/source/op/pt/dpa4c/graph_compress.cu
@@ -2,29 +2,41 @@
 //
 // Torch bindings of the compressed degree-wise DPA4C descriptor.
 //
-// This translation unit owns argument validation, the scalar-width dispatch,
-// and the operator registration. The kernels themselves are instantiated in
-// one translation unit per scalar width.
+// This translation unit owns argument validation, the dispatch over the scalar
+// width and the angular degree, the node-type scan, the prologue, epilogue and
+// per-node kernels of the fused energy-force operator, and the operator
+// registration. The descriptor kernels are instantiated in the generated
+// shards of shard.cu.in: one per angular degree for the widest scalar widths
+// and one per width for the others.
 
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
 
 #include 
+#include 
 #include 
 #include 
+#include 
 #include 
 #include 
 
+#include "../canonical_index.h"
 #include "../graph_ops.h"
+#include "../long_range/charge_sum.h"
 #include "graph_compress_launch.h"
 
 namespace {
 
 using deepmd_dpa4c::Arguments;
 using deepmd_dpa4c::IndexKind;
+using deepmd_dpa4c::kFoldedModes;
+using deepmd_dpa4c::kForceRangeWord;
+using deepmd_dpa4c::kForceWords;
+using deepmd_dpa4c::kTypeWords;
 using deepmd_dpa4c::Profile;
 
 #define DPA4C_CHECK_LAUNCH(name)                                              \
@@ -40,13 +52,15 @@ struct Dimensions {
   int degree_one;
   int coupling_records;
   int spin_channels;
+  int basis;
 };
 
 template 
 Dimensions dimensions_of() {
   using P = Profile;
-  return {P::MomentWidth, P::OutputWidth, P::C1,
-          deepmd_dpa4c::coupling_record_count(Lmax), P::Cs};
+  return {P::MomentWidth, P::OutputWidth,
+          P::C1,          deepmd_dpa4c::coupling_record_count(Lmax),
+          P::Cs,          P::Basis};
 }
 
 template 
@@ -63,6 +77,14 @@ Dimensions dimensions_for(int lmax) {
   }
 }
 
+/// Whether the edge scans prefetch the bond vectors of a canonical graph whose
+/// edge storage holds `bytes` of them: when they exceed the L2 cache of the
+/// current device, as those of a graph that fills the storage then stream from
+/// DRAM.
+bool prefetches_bonds(long bytes) {
+  return bytes > at::cuda::getCurrentDeviceProperties()->l2CacheSize;
+}
+
 Dimensions profile_dimensions(int channels, int lmax, bool has_spin) {
 #define DPA4C_DIMENSIONS(width)                           \
   if (channels == width) {                                \
@@ -74,19 +96,48 @@ Dimensions profile_dimensions(int channels, int lmax, bool has_spin) {
   TORCH_CHECK(false, "dpa4c_graph_compress: unsupported channels ", channels);
 }
 
+template 
+void launch(bool backward, const Arguments& arguments, cudaStream_t stream) {
+  if (backward) {
+    deepmd_dpa4c::launch_backward(arguments, stream);
+  } else {
+    deepmd_dpa4c::launch_forward(arguments, stream);
+  }
+}
+
+// The arguments are validated before they reach this dispatch. The default is
+// still checked rather than folded into the highest degree, so that a degree
+// outside the compiled set fails loudly instead of running a kernel for a
+// different model.
+template 
+void dispatch_degree(bool backward,
+                     const Arguments& arguments,
+                     cudaStream_t stream) {
+  switch (arguments.lmax) {
+    case 2:
+      launch(backward, arguments, stream);
+      return;
+    case 3:
+      launch(backward, arguments, stream);
+      return;
+    case 4:
+      launch(backward, arguments, stream);
+      return;
+    default:
+      TORCH_CHECK(false, "dpa4c_graph_compress: uncompiled angular degree ",
+                  arguments.lmax);
+  }
+}
+
 void dispatch(int channels,
               bool backward,
               const Arguments& arguments,
               cudaStream_t stream) {
-#define DPA4C_DISPATCH(width)                                    \
-  if (channels == width) {                                       \
-    if (backward) {                                              \
-      deepmd_dpa4c::launch_backward_c##width(arguments, stream); \
-    } else {                                                     \
-      deepmd_dpa4c::launch_forward_c##width(arguments, stream);  \
-    }                                                            \
-    DPA4C_CHECK_LAUNCH("dpa4c_graph_compress");                  \
-    return;                                                      \
+#define DPA4C_DISPATCH(width)                            \
+  if (channels == width) {                               \
+    dispatch_degree(backward, arguments, stream); \
+    DPA4C_CHECK_LAUNCH("dpa4c_graph_compress");          \
+    return;                                              \
   }
   DPA4C_FOR_EACH_CHANNEL(DPA4C_DISPATCH)
 #undef DPA4C_DISPATCH
@@ -101,6 +152,10 @@ struct Payload {
   torch::Tensor destination_row_ptr;
   torch::Tensor atype;
   torch::Tensor table;
+  // Low-rank table of the single-type and slotted scans and its expansion,
+  // empty together when the descriptor carries none.
+  torch::Tensor basis_table;
+  torch::Tensor basis_matrix;
   torch::Tensor pair_film;
   torch::Tensor pair_mixing;
   torch::Tensor type_embedding;
@@ -160,11 +215,12 @@ Arguments build_arguments(const Payload& payload,
   for (const torch::Tensor* tensor :
        {&edge_vec, &payload.edge_index, &payload.edge_mask,
         &payload.destination_order, &payload.destination_row_ptr,
-        &payload.atype, &payload.table, &payload.pair_film,
-        &payload.pair_mixing, &payload.type_embedding,
-        &payload.readout_matrices, &payload.coupling_meta,
-        &payload.coupling_entry, &payload.coupling_value, &payload.output_mean,
-        &payload.output_inv_std, &payload.contact_radius}) {
+        &payload.atype, &payload.table, &payload.basis_table,
+        &payload.basis_matrix, &payload.pair_film, &payload.pair_mixing,
+        &payload.type_embedding, &payload.readout_matrices,
+        &payload.coupling_meta, &payload.coupling_entry,
+        &payload.coupling_value, &payload.output_mean, &payload.output_inv_std,
+        &payload.contact_radius}) {
     TORCH_CHECK(tensor->is_cuda() && tensor->device() == device,
                 "dpa4c_graph_compress: every tensor input must be a CUDA "
                 "tensor on the device of edge_vec");
@@ -185,21 +241,26 @@ Arguments build_arguments(const Payload& payload,
       payload.destination_row_ptr.scalar_type() == torch::kInt64 &&
           payload.atype.scalar_type() == torch::kInt64,
       "dpa4c_graph_compress: row pointers and atom types must be int64");
+  TORCH_CHECK(payload.atype.dim() == 1,
+              "dpa4c_graph_compress: atype must be one-dimensional");
   TORCH_CHECK(payload.coupling_meta.scalar_type() == torch::kInt32 &&
                   payload.coupling_entry.scalar_type() == torch::kInt32,
               "dpa4c_graph_compress: the coupling layout must be int32");
   for (const torch::Tensor* tensor :
-       {&payload.table, &payload.pair_film, &payload.pair_mixing,
-        &payload.type_embedding, &payload.readout_matrices,
-        &payload.coupling_value, &payload.output_mean, &payload.output_inv_std,
+       {&payload.table, &payload.basis_table, &payload.basis_matrix,
+        &payload.pair_film, &payload.pair_mixing, &payload.type_embedding,
+        &payload.readout_matrices, &payload.coupling_value,
+        &payload.output_mean, &payload.output_inv_std,
         &payload.contact_radius}) {
     TORCH_CHECK(tensor->scalar_type() == torch::kFloat32,
                 "dpa4c_graph_compress: tables and weights must be fp32");
   }
   TORCH_CHECK(payload.lmax >= 2 && payload.lmax <= 4,
               "dpa4c_graph_compress: lmax must be 2, 3, or 4");
-  TORCH_CHECK(payload.table_max >= payload.rcut,
-              "dpa4c_graph_compress: table_max must cover rcut");
+  TORCH_CHECK(std::isfinite(payload.table_stride) && payload.table_stride > 0.0,
+              "dpa4c_graph_compress: table_stride must be positive and "
+              "finite, got ",
+              payload.table_stride);
   TORCH_CHECK(payload.f_inner >= 0.0 && payload.f_outer >= payload.f_inner,
               "dpa4c_graph_compress: the bridging switch needs 0 <= f_inner <= "
               "f_outer, got ",
@@ -215,6 +276,40 @@ Arguments build_arguments(const Payload& payload,
   TORCH_CHECK(payload.table.dim() == 2 && payload.table.size(0) > 0 &&
                   payload.table.size(1) == 6 * (channels + radial_modes),
               "dpa4c_graph_compress: invalid radial table shape");
+  // The scans bound the bond-length windows of a row by the table reach, so
+  // the reach lies within the knots of the table.
+  const double table_extent =
+      static_cast(payload.table.size(0)) * payload.table_stride;
+  TORCH_CHECK(payload.table_max >= payload.rcut &&
+                  payload.table_max <= table_extent * (1.0 + 1e-12),
+              "dpa4c_graph_compress: table_max must cover rcut within the ",
+              payload.table.size(0), " intervals of the radial table, got ",
+              payload.table_max, " for rcut ", payload.rcut, " and stride ",
+              payload.table_stride);
+  // A low-rank table shares the knots of the radial table and holds the
+  // degree-one channels, the basis functions and the modes.
+  if (payload.basis_table.numel() != 0) {
+    TORCH_CHECK(widths.basis > 0, "dpa4c_graph_compress: ",
+                payload.spin.dim() == 2
+                    ? std::string("a spin-conditioned descriptor admits")
+                    : std::to_string(channels) + " channels admit",
+                " no low-rank radial table");
+    TORCH_CHECK(payload.basis_table.dim() == 2 &&
+                    payload.basis_table.size(0) == payload.table.size(0) &&
+                    payload.basis_table.size(1) ==
+                        6 * (widths.degree_one + widths.basis + radial_modes),
+                "dpa4c_graph_compress: invalid low-rank radial table shape");
+    TORCH_CHECK(
+        payload.basis_matrix.sizes() ==
+            torch::IntArrayRef({widths.basis, channels - widths.degree_one}),
+        "dpa4c_graph_compress: the low-rank expansion must have "
+        "shape (",
+        widths.basis, ", ", channels - widths.degree_one, ")");
+  } else {
+    TORCH_CHECK(payload.basis_matrix.numel() == 0,
+                "dpa4c_graph_compress: a low-rank expansion requires its "
+                "table");
+  }
   TORCH_CHECK(payload.type_embedding.dim() == 2 && type_count > 1,
               "dpa4c_graph_compress: invalid type embedding shape");
   TORCH_CHECK(
@@ -248,6 +343,22 @@ Arguments build_arguments(const Payload& payload,
                   torch::IntArrayRef({static_cast(type_count)}),
               "dpa4c_graph_compress: contact_radius must have shape "
               "(type_count,)");
+  // The scans read the radial tables, the low-rank expansion, the mode mixing
+  // and the per-type spin table as float4, and the PairFiLM and ordered spin
+  // caches as float2.
+  const auto aligned = [](const torch::Tensor& tensor, std::uintptr_t bytes) {
+    return tensor.numel() == 0 ||
+           reinterpret_cast(tensor.data_ptr()) % bytes == 0;
+  };
+  TORCH_CHECK(
+      aligned(payload.table, 16) && aligned(payload.basis_table, 16) &&
+          aligned(payload.basis_matrix, 16) &&
+          aligned(payload.pair_mixing, 16) && aligned(payload.spin_type, 16) &&
+          aligned(payload.pair_film, 8) && aligned(payload.spin_pair, 8),
+      "dpa4c_graph_compress: the radial tables, the low-rank "
+      "expansion, the mode mixing and the per-type spin table must be "
+      "16-byte aligned, and the PairFiLM and ordered spin caches "
+      "8-byte aligned");
   // A run covers ``node_count`` destination rows starting at ``node_begin``;
   // the type table is system-wide because neighbor lookups index it with
   // absolute source indices. Whole-system entry points additionally require
@@ -297,6 +408,7 @@ Arguments build_arguments(const Payload& payload,
   arguments.spin_type =
       has_spin ? payload.spin_type.data_ptr() : nullptr;
   arguments.node_count = node_count;
+  arguments.system_node_count = node_count;
   arguments.edge_count = edge_vec.size(0);
   arguments.lmax = static_cast(payload.lmax);
   arguments.interval_count = static_cast(payload.table.size(0));
@@ -317,6 +429,9 @@ Arguments build_arguments(const Payload& payload,
     arguments.bridged = true;
   }
   arguments.canonical = payload.canonical;
+  arguments.prefetch_bonds =
+      payload.canonical &&
+      prefetches_bonds(edge_vec.size(0) * 3 * static_cast(sizeof(float)));
   arguments.index_kind = index_kind_of(payload.edge_index);
   arguments.edge_index = payload.edge_index.data_ptr();
   arguments.destination_order = payload.destination_order.numel() != 0
@@ -327,6 +442,10 @@ Arguments build_arguments(const Payload& payload,
                             ? payload.edge_mask.data_ptr()
                             : nullptr;
   arguments.destination_row_ptr = payload.destination_row_ptr.data_ptr();
+  if (payload.basis_table.numel() != 0) {
+    arguments.basis_table = payload.basis_table.data_ptr();
+    arguments.basis_matrix = payload.basis_matrix.data_ptr();
+  }
   arguments.node_begin = payload.node_begin;
   arguments.atype = payload.atype.data_ptr();
   arguments.table = payload.table.data_ptr();
@@ -373,6 +492,138 @@ void attach_pair_term(Arguments& arguments,
   arguments.bridged = true;
 }
 
+/// Words that one scan of the node types of a canonical system writes,
+/// allocated cleared, for the edge kernels that follow it: whether any node
+/// carries a type other than that of node zero, which selects the folded scans
+/// (`Arguments::type_mismatch`), the largest length-scale radius among the
+/// node types, which bounds the bridging windows of every row
+/// (`Arguments::largest_radius`), and the bitmap of the node types, which
+/// numbers the types of a system for the slotted scans
+/// (`Arguments::type_present`).
+struct NodeTypeWords {
+  int type_mismatch;
+  float largest_radius;
+  unsigned type_present[kTypeWords];
+};
+
+/// Set the bit of `type` in a bitmap of node types. A negative type, and a
+/// type beyond the words of the bitmap, which a launch of more types never
+/// reads, set none.
+__device__ __forceinline__ void mark_type(long type,
+                                          unsigned (&present)[kTypeWords]) {
+#pragma unroll
+  for (int word = 0; word < kTypeWords; ++word) {
+    if (type >> 5 == word) {
+      present[word] |= 1u << (type & 31);
+    }
+  }
+}
+
+/// Fold the node types that the threads of a block have scanned into `words`.
+/// Every thread of the block takes part: `same` tells whether all of the
+/// thread's nodes carry the type of node zero, `largest` is the largest
+/// radius among them, and `present` the bitmap of their types. Radii are
+/// positive, so their bit patterns order as the integers do and the integer
+/// maximum is the floating-point one.
+__device__ __forceinline__ void publish_node_types(
+    bool same,
+    float largest,
+    const unsigned (&present)[kTypeWords],
+    NodeTypeWords* words) {
+  constexpr int kLanes = 32;
+#pragma unroll
+  for (int offset = kLanes / 2; offset > 0; offset >>= 1) {
+    largest = fmaxf(largest, __shfl_xor_sync(0xffffffffu, largest, offset));
+  }
+  if (threadIdx.x % kLanes == 0 && largest > 0.0f) {
+    atomicMax(reinterpret_cast(&words->largest_radius),
+              __float_as_int(largest));
+  }
+#pragma unroll
+  for (int word = 0; word < kTypeWords; ++word) {
+    unsigned bits = present[word];
+#pragma unroll
+    for (int offset = kLanes / 2; offset > 0; offset >>= 1) {
+      bits |= __shfl_xor_sync(0xffffffffu, bits, offset);
+    }
+    if (threadIdx.x % kLanes == 0 && bits != 0) {
+      atomicOr(&words->type_present[word], bits);
+    }
+  }
+  if (!__syncthreads_and(same) && threadIdx.x == 0) {
+    atomicOr(&words->type_mismatch, 1);
+  }
+}
+
+// Scan the node types of a canonical system into `words`; a null
+// `contact_radius` leaves the largest radius at zero.
+__global__ void node_type_kernel(long count,
+                                 const long* __restrict__ atype,
+                                 const float* __restrict__ contact_radius,
+                                 NodeTypeWords* __restrict__ words) {
+  const long first = atype[0];
+  bool same = true;
+  float largest = 0.0f;
+  unsigned present[kTypeWords] = {};
+  for (long node = blockIdx.x * static_cast(blockDim.x) + threadIdx.x;
+       node < count; node += static_cast(blockDim.x) * gridDim.x) {
+    const long type = atype[node];
+    same &= type == first;
+    mark_type(type, present);
+    if (contact_radius != nullptr && type >= 0) {
+      largest = fmaxf(largest, __ldg(contact_radius + type));
+    }
+  }
+  publish_node_types(same, largest, present, words);
+}
+
+/// Radius table of the node-type scan: that of a launch with a bridging
+/// window, whose edge scans bound the window by the largest radius, and none
+/// otherwise.
+const float* windowed_radii(const Arguments& arguments) {
+  return arguments.f_outer > arguments.f_inner ? arguments.contact_radius
+                                               : nullptr;
+}
+
+/// Bitmap of the node types of `words` for a launch that its slotted scans
+/// may serve: one whose types, the padding type included, the bitmap covers,
+/// of at most `kFoldedModes` radial modes, whose descriptor carries modes or
+/// a low-rank table for the slots to hold. Null for every other launch, whose
+/// systems the other kernels of the profile scan.
+const unsigned* present_types(const Arguments& arguments,
+                              const NodeTypeWords* words) {
+  const bool slotted_features =
+      arguments.radial_modes > 0 || arguments.basis_table != nullptr;
+  return slotted_features && arguments.type_count <= 32 * kTypeWords &&
+                 arguments.radial_modes <= kFoldedModes
+             ? words->type_present
+             : nullptr;
+}
+
+/// Scan the node types of a canonical system (see `NodeTypeWords`) and point
+/// the launch at the words.
+torch::Tensor scan_node_types(const torch::Tensor& atype,
+                              Arguments& arguments,
+                              cudaStream_t stream) {
+  auto words =
+      torch::zeros({static_cast(sizeof(NodeTypeWords) / sizeof(int))},
+                   atype.options().dtype(torch::kInt32));
+  auto* node_types = reinterpret_cast(words.data_ptr());
+  const long count = atype.numel();
+  if (count > 0) {
+    constexpr int kThreads = 256;
+    const int blocks = static_cast(
+        std::min((count + kThreads - 1) / kThreads, 1024));
+    node_type_kernel<<>>(
+        count, atype.data_ptr(), windowed_radii(arguments), node_types);
+    DPA4C_CHECK_LAUNCH("dpa4c node types");
+  }
+  arguments.type_mismatch = &node_types->type_mismatch;
+  arguments.largest_radius = &node_types->largest_radius;
+  arguments.type_present = present_types(arguments, node_types);
+  return words;
+}
+
 }  // namespace
 
 std::tuple dpa4c_graph_compress(
@@ -383,6 +634,8 @@ std::tuple dpa4c_graph_compress(
     torch::Tensor destination_row_ptr,
     torch::Tensor atype,
     torch::Tensor table,
+    torch::Tensor basis_table,
+    torch::Tensor basis_matrix,
     torch::Tensor pair_film,
     torch::Tensor pair_mixing,
     torch::Tensor type_embedding,
@@ -405,34 +658,16 @@ std::tuple dpa4c_graph_compress(
     double degree_floor,
     double f_inner,
     double f_outer) {
-  const Payload payload{edge_index,
-                        edge_mask,
-                        destination_order,
-                        destination_row_ptr,
-                        atype,
-                        table,
-                        pair_film,
-                        pair_mixing,
-                        type_embedding,
-                        readout_matrices,
-                        coupling_meta,
-                        coupling_entry,
-                        coupling_value,
-                        output_mean,
-                        output_inv_std,
-                        spin,
-                        spin_pair,
-                        spin_type,
-                        contact_radius,
-                        canonical,
-                        lmax,
-                        table_stride,
-                        table_max,
-                        rcut,
-                        eps,
-                        degree_floor,
-                        f_inner,
-                        f_outer};
+  const Payload payload{edge_index,          edge_mask,      destination_order,
+                        destination_row_ptr, atype,          table,
+                        basis_table,         basis_matrix,   pair_film,
+                        pair_mixing,         type_embedding, readout_matrices,
+                        coupling_meta,       coupling_entry, coupling_value,
+                        output_mean,         output_inv_std, spin,
+                        spin_pair,           spin_type,      contact_radius,
+                        canonical,           lmax,           table_stride,
+                        table_max,           rcut,           eps,
+                        degree_floor,        f_inner,        f_outer};
   const long node_count = destination_row_ptr.numel() - 1;
   // The destination row pointer defines the node axis; the type table must
   // describe exactly that axis, or the two disagree on how many nodes exist.
@@ -452,11 +687,14 @@ std::tuple dpa4c_graph_compress(
     return {descriptor, state};
   }
   auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous();
+  const auto stream = at::cuda::getCurrentCUDAStream();
   Arguments arguments =
       build_arguments(payload, edge_vec_float, channels, widths);
   arguments.descriptor = descriptor.data_ptr();
   arguments.state_out = state.data_ptr();
-  dispatch(channels, false, arguments, at::cuda::getCurrentCUDAStream());
+  const torch::Tensor node_types =
+      canonical ? scan_node_types(atype, arguments, stream) : torch::Tensor();
+  dispatch(channels, false, arguments, stream);
   return {descriptor, state};
 }
 
@@ -470,6 +708,8 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient,
                                    torch::Tensor destination_row_ptr,
                                    torch::Tensor atype,
                                    torch::Tensor table,
+                                   torch::Tensor basis_table,
+                                   torch::Tensor basis_matrix,
                                    torch::Tensor pair_film,
                                    torch::Tensor pair_mixing,
                                    torch::Tensor type_embedding,
@@ -495,34 +735,16 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient,
                                    torch::Tensor pair_table,
                                    torch::Tensor pair_seed,
                                    bool reuse_state) {
-  const Payload payload{edge_index,
-                        edge_mask,
-                        destination_order,
-                        destination_row_ptr,
-                        atype,
-                        table,
-                        pair_film,
-                        pair_mixing,
-                        type_embedding,
-                        readout_matrices,
-                        coupling_meta,
-                        coupling_entry,
-                        coupling_value,
-                        output_mean,
-                        output_inv_std,
-                        spin,
-                        spin_pair,
-                        spin_type,
-                        contact_radius,
-                        canonical,
-                        lmax,
-                        table_stride,
-                        table_max,
-                        rcut,
-                        eps,
-                        degree_floor,
-                        f_inner,
-                        f_outer};
+  const Payload payload{edge_index,          edge_mask,      destination_order,
+                        destination_row_ptr, atype,          table,
+                        basis_table,         basis_matrix,   pair_film,
+                        pair_mixing,         type_embedding, readout_matrices,
+                        coupling_meta,       coupling_entry, coupling_value,
+                        output_mean,         output_inv_std, spin,
+                        spin_pair,           spin_type,      contact_radius,
+                        canonical,           lmax,           table_stride,
+                        table_max,           rcut,           eps,
+                        degree_floor,        f_inner,        f_outer};
   const bool has_spin = spin.dim() == 2;
   const long node_count = destination_row_ptr.numel() - 1;
   // The destination row pointer defines the node axis; the type table must
@@ -608,7 +830,10 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient,
                      pair_seed_c.data_ptr(),
                      pair_energy.data_ptr());
   }
-  dispatch(channels, true, arguments, at::cuda::getCurrentCUDAStream());
+  const auto stream = at::cuda::getCurrentCUDAStream();
+  const torch::Tensor node_types =
+      canonical ? scan_node_types(atype, arguments, stream) : torch::Tensor();
+  dispatch(channels, true, arguments, stream);
   return {edge_gradient.to(edge_vec.scalar_type()), spin_gradient,
           edge_spin_gradient, pair_energy};
 }
@@ -623,6 +848,8 @@ dpa4c_graph_compress_backward(torch::Tensor descriptor_gradient,
                               torch::Tensor destination_row_ptr,
                               torch::Tensor atype,
                               torch::Tensor table,
+                              torch::Tensor basis_table,
+                              torch::Tensor basis_matrix,
                               torch::Tensor pair_film,
                               torch::Tensor pair_mixing,
                               torch::Tensor type_embedding,
@@ -649,12 +876,12 @@ dpa4c_graph_compress_backward(torch::Tensor descriptor_gradient,
                               torch::Tensor pair_seed) {
   return dpa4c_graph_compress_backward_impl(
       descriptor_gradient, state, edge_vec, edge_index, edge_mask,
-      destination_order, destination_row_ptr, atype, table, pair_film,
-      pair_mixing, type_embedding, readout_matrices, coupling_meta,
-      coupling_entry, coupling_value, output_mean, output_inv_std, spin,
-      spin_pair, spin_type, contact_radius, canonical, lmax, table_stride,
-      table_max, rcut, eps, degree_floor, f_inner, f_outer, pair_table,
-      pair_seed, false);
+      destination_order, destination_row_ptr, atype, table, basis_table,
+      basis_matrix, pair_film, pair_mixing, type_embedding, readout_matrices,
+      coupling_meta, coupling_entry, coupling_value, output_mean,
+      output_inv_std, spin, spin_pair, spin_type, contact_radius, canonical,
+      lmax, table_stride, table_max, rcut, eps, degree_floor, f_inner, f_outer,
+      pair_table, pair_seed, false);
 }
 
 std::tuple dpa4c_canonical_compress(
@@ -663,6 +890,8 @@ std::tuple dpa4c_canonical_compress(
     torch::Tensor destination_row_ptr,
     torch::Tensor atype,
     torch::Tensor table,
+    torch::Tensor basis_table,
+    torch::Tensor basis_matrix,
     torch::Tensor pair_film,
     torch::Tensor pair_mixing,
     torch::Tensor type_embedding,
@@ -684,6 +913,7 @@ std::tuple dpa4c_canonical_compress(
     double degree_floor,
     double f_inner,
     double f_outer) {
+  source = deepmd::canonical_index_view(source);
   TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0),
               "dpa4c_canonical_compress: source and edge_vec must share the "
               "edge axis");
@@ -691,10 +921,11 @@ std::tuple dpa4c_canonical_compress(
   auto destination_order = torch::empty({0}, source.options());
   return dpa4c_graph_compress(
       edge_vec, source, edge_mask, destination_order, destination_row_ptr,
-      atype, table, pair_film, pair_mixing, type_embedding, readout_matrices,
-      coupling_meta, coupling_entry, coupling_value, output_mean,
-      output_inv_std, spin, spin_pair, spin_type, contact_radius, true, lmax,
-      table_stride, table_max, rcut, eps, degree_floor, f_inner, f_outer);
+      atype, table, basis_table, basis_matrix, pair_film, pair_mixing,
+      type_embedding, readout_matrices, coupling_meta, coupling_entry,
+      coupling_value, output_mean, output_inv_std, spin, spin_pair, spin_type,
+      contact_radius, true, lmax, table_stride, table_max, rcut, eps,
+      degree_floor, f_inner, f_outer);
 }
 
 std::tuple
@@ -705,6 +936,8 @@ dpa4c_canonical_compress_backward_common(torch::Tensor descriptor_gradient,
                                          torch::Tensor destination_row_ptr,
                                          torch::Tensor atype,
                                          torch::Tensor table,
+                                         torch::Tensor basis_table,
+                                         torch::Tensor basis_matrix,
                                          torch::Tensor pair_film,
                                          torch::Tensor pair_mixing,
                                          torch::Tensor type_embedding,
@@ -727,23 +960,24 @@ dpa4c_canonical_compress_backward_common(torch::Tensor descriptor_gradient,
                                          double f_inner,
                                          double f_outer,
                                          bool reuse_state) {
+  source = deepmd::canonical_index_view(source);
   TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0),
               "dpa4c_canonical_compress_backward: source and edge_vec must "
               "share the edge axis");
   auto edge_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool));
   auto destination_order = torch::empty({0}, source.options());
   // The compact descriptor operators carry no analytical pair potential; the
-  // fused energy-gradient operator evaluates it on this ABI.
+  // fused energy-force operator evaluates it on this ABI.
   auto no_pair = torch::empty({0}, edge_vec.options().dtype(torch::kFloat32));
   auto [edge_gradient, spin_gradient, edge_spin_gradient, pair_energy] =
       dpa4c_graph_compress_backward_impl(
           descriptor_gradient, state, edge_vec, source, edge_mask,
-          destination_order, destination_row_ptr, atype, table, pair_film,
-          pair_mixing, type_embedding, readout_matrices, coupling_meta,
-          coupling_entry, coupling_value, output_mean, output_inv_std, spin,
-          spin_pair, spin_type, contact_radius, true, lmax, table_stride,
-          table_max, rcut, eps, degree_floor, f_inner, f_outer, no_pair,
-          no_pair, reuse_state);
+          destination_order, destination_row_ptr, atype, table, basis_table,
+          basis_matrix, pair_film, pair_mixing, type_embedding,
+          readout_matrices, coupling_meta, coupling_entry, coupling_value,
+          output_mean, output_inv_std, spin, spin_pair, spin_type,
+          contact_radius, true, lmax, table_stride, table_max, rcut, eps,
+          degree_floor, f_inner, f_outer, no_pair, no_pair, reuse_state);
   return {edge_gradient, spin_gradient, edge_spin_gradient};
 }
 
@@ -755,6 +989,8 @@ dpa4c_canonical_compress_backward(torch::Tensor descriptor_gradient,
                                   torch::Tensor destination_row_ptr,
                                   torch::Tensor atype,
                                   torch::Tensor table,
+                                  torch::Tensor basis_table,
+                                  torch::Tensor basis_matrix,
                                   torch::Tensor pair_film,
                                   torch::Tensor pair_mixing,
                                   torch::Tensor type_embedding,
@@ -778,10 +1014,10 @@ dpa4c_canonical_compress_backward(torch::Tensor descriptor_gradient,
                                   double f_outer) {
   return dpa4c_canonical_compress_backward_common(
       descriptor_gradient, state, edge_vec, source, destination_row_ptr, atype,
-      table, pair_film, pair_mixing, type_embedding, readout_matrices,
-      coupling_meta, coupling_entry, coupling_value, output_mean,
-      output_inv_std, spin, spin_pair, spin_type, contact_radius, lmax,
-      table_stride, table_max, rcut, eps, degree_floor, f_inner, f_outer,
+      table, basis_table, basis_matrix, pair_film, pair_mixing, type_embedding,
+      readout_matrices, coupling_meta, coupling_entry, coupling_value,
+      output_mean, output_inv_std, spin, spin_pair, spin_type, contact_radius,
+      lmax, table_stride, table_max, rcut, eps, degree_floor, f_inner, f_outer,
       false);
 }
 
@@ -793,6 +1029,8 @@ dpa4c_canonical_compress_backward_inplace(torch::Tensor descriptor_gradient,
                                           torch::Tensor destination_row_ptr,
                                           torch::Tensor atype,
                                           torch::Tensor table,
+                                          torch::Tensor basis_table,
+                                          torch::Tensor basis_matrix,
                                           torch::Tensor pair_film,
                                           torch::Tensor pair_mixing,
                                           torch::Tensor type_embedding,
@@ -816,126 +1054,897 @@ dpa4c_canonical_compress_backward_inplace(torch::Tensor descriptor_gradient,
                                           double f_outer) {
   return dpa4c_canonical_compress_backward_common(
       descriptor_gradient, state, edge_vec, source, destination_row_ptr, atype,
-      table, pair_film, pair_mixing, type_embedding, readout_matrices,
-      coupling_meta, coupling_entry, coupling_value, output_mean,
-      output_inv_std, spin, spin_pair, spin_type, contact_radius, lmax,
-      table_stride, table_max, rcut, eps, degree_floor, f_inner, f_outer, true);
+      table, basis_table, basis_matrix, pair_film, pair_mixing, type_embedding,
+      readout_matrices, coupling_meta, coupling_entry, coupling_value,
+      output_mean, output_inv_std, spin, spin_pair, spin_type, contact_radius,
+      lmax, table_stride, table_max, rcut, eps, degree_floor, f_inner, f_outer,
+      true);
 }
 
-// Energy and edge cotangent of one compressed inference step, evaluated over
+namespace {
+
+// Threads of one block of the prologue and of the per-node results of a
+// batch, and the logical threads whose nodes one block of the single-frame
+// epilogue reduces. The epilogue reduces its totals in double precision,
+// which most devices without a data-centre FP64 rate execute at a small
+// fraction of the single-precision rate, so its node slices are kept to two
+// warps and spread over as many multiprocessors as the node count allows.
+constexpr int kEpilogueThreads = 64;
+// Blocks of the epilogue reduction. The count follows the node count alone,
+// so the summation order, and with it every rounding, is the same on every
+// device.
+constexpr long kEpilogueBlocks = 512;
+
+/// Words of one evaluation, allocated cleared: those of the scan of the node
+/// types, which the prologue performs, and the arrivals counted by the
+/// last-block reduction of the epilogue.
+struct EvaluationWords {
+  NodeTypeWords node_types;
+  unsigned int arrivals;
+};
+
+__device__ __forceinline__ double fixed_point_value(unsigned long long word) {
+  return static_cast(static_cast(word)) *
+         (1.0 / static_cast(deepmd_dpa4c::kForceFixedScale));
+}
+
+// Per-node preparation of one evaluation: the node types with virtual atoms
+// (negative types) mapped onto type zero, the ownership mask that selects the
+// energy-contributing nodes -- a real atom among the leading ``n_local``
+// nodes of its frame -- as the int32 output mask, as the fp64 cotangent that
+// seeds the fitting backward and, unless ``charged`` is null, as the boolean
+// mask of the nodes that carry the charges of a long-range module, and
+// cleared fixed-point force totals. The frame of a node comes from
+// ``frame_row_ptr`` when the batch holds several frames and is frame zero
+// otherwise; nodes past the last frame are padding, which no frame owns.
+__global__ void energy_force_prologue_kernel(
+    long node_count,
+    long frame_count,
+    const long* __restrict__ n_local,
+    const long* __restrict__ frame_row_ptr,
+    const long* __restrict__ atype,
+    long* __restrict__ node_type,
+    double* __restrict__ seed,
+    int* __restrict__ mask,
+    bool* __restrict__ charged,
+    unsigned long long* __restrict__ force,
+    unsigned long long* __restrict__ spin_force,
+    const float* __restrict__ contact_radius,
+    EvaluationWords* __restrict__ words) {
+  const long first = atype[0] < 0 ? 0 : atype[0];
+  bool same = true;
+  float largest = 0.0f;
+  unsigned present[kTypeWords] = {};
+  for (long node = blockIdx.x * static_cast(blockDim.x) + threadIdx.x;
+       node < node_count; node += static_cast(blockDim.x) * gridDim.x) {
+    const long raw = atype[node];
+    const long type = raw < 0 ? 0 : raw;
+    node_type[node] = type;
+    same &= type == first;
+    mark_type(type, present);
+    if (contact_radius != nullptr) {
+      largest = fmaxf(largest, __ldg(contact_radius + type));
+    }
+    long frame = 0;
+    long frame_begin = 0;
+    if (frame_count > 1) {
+      // Last frame whose span begins at or before the node.
+      long low = 0;
+      long high = frame_count - 1;
+      while (low < high) {
+        const long middle = (low + high + 1) / 2;
+        if (frame_row_ptr[middle] <= node) {
+          low = middle;
+        } else {
+          high = middle - 1;
+        }
+      }
+      frame = low;
+      frame_begin = frame_row_ptr[frame];
+    }
+    const bool owned = raw >= 0 && node - frame_begin < n_local[frame];
+    seed[node] = owned ? 1.0 : 0.0;
+    mask[node] = owned ? 1 : 0;
+    if (charged != nullptr) {
+      charged[node] = owned;
+    }
+#pragma unroll
+    for (int word = 0; word < kForceWords; ++word) {
+      force[node * kForceWords + word] = 0ull;
+      if (spin_force != nullptr) {
+        spin_force[node * kForceWords + word] = 0ull;
+      }
+    }
+  }
+  publish_node_types(same, largest, present, &words->node_types);
+}
+
+// Frame totals the single-frame epilogue reduces: the energy and the nine
+// virial components.
+constexpr int kEpilogueTotals = 10;
+
+// Per-node results of one evaluation of a single frame, and its totals: the
+// fixed-point force (and magnetic force) totals in fp64, the node energies
+// masked by ownership, zeros for a per-atom virial nobody asked for, and the
+// frame energy and virial. Each block reduces a fixed slice of the nodes in a
+// fixed order, the last block to finish folds the block partials in a fixed
+// order, and only the nodes of the frame, ``[0, n_node[0])``, contribute.
+//
+// The slice of a block is that of kEpilogueThreads logical threads, and every
+// total of a logical thread is carried by a warp of its own: thread ``t`` of
+// the warps of total ``k`` accumulates total ``k`` of the nodes of logical
+// thread ``t`` in node order, each warp reduces its 32 values by the same
+// butterfly, and the warp sums and the block partials combine in the same
+// order whichever warp holds them. The order of every sum is therefore the
+// one of a block of kEpilogueThreads threads carrying all totals, while the
+// double-precision work, which most devices run at a small fraction of the
+// single-precision rate, spreads over kEpilogueTotals times as many warps.
+// The warps of a total also write the per-node results that pair with it.
+//
+// A long-range module adds its node energies, its forces at fixed charges and
+// its per-node virials, ``extra_energy``, ``extra_force`` and
+// ``extra_virial``, onto those of the descriptor and the fitting; all three
+// are null for a model without one.
+__global__ void __launch_bounds__(kEpilogueThreads* kEpilogueTotals)
+    energy_force_epilogue_kernel(
+        long node_count,
+        const long* __restrict__ n_node,
+        const unsigned long long* __restrict__ force_fixed,
+        const unsigned long long* __restrict__ spin_fixed,
+        const float* __restrict__ spin_gradient,
+        const double* __restrict__ raw_energy,
+        const double* __restrict__ seed,
+        const float* __restrict__ node_virial,
+        const double* __restrict__ extra_energy,
+        const float* __restrict__ extra_force,
+        const float* __restrict__ extra_virial,
+        double* __restrict__ force,
+        double* __restrict__ force_mag,
+        double* __restrict__ atom_energy,
+        float* __restrict__ atom_virial_zero,
+        double* __restrict__ partial,
+        EvaluationWords* __restrict__ words,
+        double* __restrict__ energy,
+        double* __restrict__ virial) {
+  constexpr int kWarps = kEpilogueThreads / 32;
+  const long frame_end = n_node[0];
+  const int total = threadIdx.x / kEpilogueThreads;
+  const int logical = threadIdx.x % kEpilogueThreads;
+  const int lane = threadIdx.x % 32;
+  const int warp = logical / 32;
+  // Total 0 is the energy and total 1 + c the virial component c; the warps
+  // of totals 1 to 3 also write force component total - 1.
+  const int component = total - 1;
+  double sum = 0.0;
+  for (long node = blockIdx.x * static_cast(kEpilogueThreads) + logical;
+       node < node_count;
+       node += static_cast(kEpilogueThreads) * gridDim.x) {
+    if (total == 0) {
+      // A node whose force contributions left the range of the fixed-point
+      // totals reports its energy as not a number, which marks the forces of
+      // its frame.
+      double value = force_fixed[node * kForceWords + kForceRangeWord] != 0
+                         ? CUDART_NAN
+                         : raw_energy[node] * seed[node];
+      if (extra_energy != nullptr) {
+        value += extra_energy[node];
+      }
+      atom_energy[node] = value;
+      if (node < frame_end) {
+        sum += value;
+      }
+      continue;
+    }
+    if (component < 3) {
+      double value =
+          fixed_point_value(force_fixed[node * kForceWords + component]);
+      if (extra_force != nullptr) {
+        value += static_cast(extra_force[node * 3 + component]);
+      }
+      force[node * 3 + component] = value;
+      if (force_mag != nullptr) {
+        // The magnetic force is minus the energy gradient with respect to the
+        // moment: its on-site part and the neighbour part of its sources.
+        force_mag[node * 3 + component] =
+            -(static_cast(spin_gradient[node * 3 + component]) +
+              fixed_point_value(spin_fixed[node * kForceWords + component]));
+      }
+    }
+    if (atom_virial_zero != nullptr) {
+      atom_virial_zero[node * 9 + component] = 0.0f;
+    }
+    if (node < frame_end) {
+      sum += static_cast(node_virial[node * 9 + component]);
+      if (extra_virial != nullptr) {
+        sum += static_cast(extra_virial[node * 9 + component]);
+      }
+    }
+  }
+
+  // === Block totals in a fixed order ===
+  __shared__ double warp_totals[kWarps][kEpilogueTotals];
+  __shared__ bool last_block;
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    sum += __shfl_down_sync(0xffffffffu, sum, offset);
+  }
+  if (lane == 0) {
+    warp_totals[warp][total] = sum;
+  }
+  __syncthreads();
+  if (threadIdx.x < kEpilogueTotals) {
+    double value = 0.0;
+    for (int w = 0; w < kWarps; ++w) {
+      value += warp_totals[w][threadIdx.x];
+    }
+    partial[blockIdx.x * kEpilogueTotals + threadIdx.x] = value;
+  }
+  // === The last block to arrive folds the block totals ===
+  __threadfence();
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    last_block = atomicAdd(&words->arrivals, 1u) == gridDim.x - 1;
+  }
+  __syncthreads();
+  if (!last_block) {
+    return;
+  }
+  __threadfence();
+  // Warp k of the block folds total k.
+  __shared__ double totals[kEpilogueTotals];
+  const int fold = threadIdx.x / 32;
+  if (fold < kEpilogueTotals) {
+    double value = 0.0;
+    for (int block = lane; block < static_cast(gridDim.x); block += 32) {
+      value += __ldcg(partial + block * kEpilogueTotals + fold);
+    }
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_down_sync(0xffffffffu, value, offset);
+    }
+    if (lane == 0) {
+      totals[fold] = value;
+      if (fold == 0) {
+        energy[0] = value;
+      } else {
+        virial[fold - 1] = value;
+      }
+    }
+  }
+  __syncthreads();
+  // A fixed-point force word cannot hold a non-finite value. A non-finite
+  // descriptor makes the energy non-finite, and a non-finite edge gradient,
+  // from which every force follows, the virial of its destination, so a
+  // frame whose totals are not finite reports every force of its nodes as
+  // not a number. Nodes past the frame belong to none and keep their forces,
+  // as in the multi-frame path.
+  bool finite = true;
+#pragma unroll
+  for (int index = 0; index < kEpilogueTotals; ++index) {
+    finite = finite && isfinite(totals[index]);
+  }
+  if (!finite) {
+    const long frame_entries =
+        3 * (frame_end < node_count ? frame_end : node_count);
+    for (long index = threadIdx.x; index < frame_entries; index += blockDim.x) {
+      force[index] = CUDART_NAN;
+      if (force_mag != nullptr) {
+        force_mag[index] = CUDART_NAN;
+      }
+    }
+  }
+}
+
+// Per-node results of a batch of several frames; the frame totals follow
+// from the segment reductions of the caller. ``extra_energy`` and
+// ``extra_force`` are the contributions of a long-range module, as in
+// energy_force_epilogue_kernel.
+__global__ void energy_force_nodes_kernel(
+    long node_count,
+    const unsigned long long* __restrict__ force_fixed,
+    const unsigned long long* __restrict__ spin_fixed,
+    const float* __restrict__ spin_gradient,
+    const double* __restrict__ raw_energy,
+    const double* __restrict__ seed,
+    const double* __restrict__ extra_energy,
+    const float* __restrict__ extra_force,
+    double* __restrict__ force,
+    double* __restrict__ force_mag,
+    double* __restrict__ atom_energy,
+    float* __restrict__ atom_virial_zero) {
+  for (long node = blockIdx.x * static_cast(blockDim.x) + threadIdx.x;
+       node < node_count; node += static_cast(blockDim.x) * gridDim.x) {
+#pragma unroll
+    for (int component = 0; component < 3; ++component) {
+      double value =
+          fixed_point_value(force_fixed[node * kForceWords + component]);
+      if (extra_force != nullptr) {
+        value += static_cast(extra_force[node * 3 + component]);
+      }
+      force[node * 3 + component] = value;
+      if (force_mag != nullptr) {
+        force_mag[node * 3 + component] =
+            -(static_cast(spin_gradient[node * 3 + component]) +
+              fixed_point_value(spin_fixed[node * kForceWords + component]));
+      }
+    }
+    double value = force_fixed[node * kForceWords + kForceRangeWord] != 0
+                       ? CUDART_NAN
+                       : raw_energy[node] * seed[node];
+    if (extra_energy != nullptr) {
+      value += extra_energy[node];
+    }
+    atom_energy[node] = value;
+    if (atom_virial_zero != nullptr) {
+#pragma unroll
+      for (int component = 0; component < 9; ++component) {
+        atom_virial_zero[node * 9 + component] = 0.0f;
+      }
+    }
+  }
+}
+
+int epilogue_blocks(long node_count) {
+  return static_cast(std::max(
+      1, std::min((node_count + kEpilogueThreads - 1) / kEpilogueThreads,
+                        kEpilogueBlocks)));
+}
+
+// Validates the graph inputs of a fused energy-force operator named ``op``
+// that its prologue reads before the runs validate the rest of the payload:
+// the frame counts, the node types and, for a windowed descriptor, the
+// contact radii, and the source CSR, which is either complete, requesting the
+// per-atom virial, or the two-entry placeholder that declines it. ``source``
+// is replaced by its typed view. Returns whether the per-atom virial is
+// requested.
+bool check_energy_force_inputs(const char* op,
+                               const torch::Tensor& edge_vec,
+                               torch::Tensor& source,
+                               const torch::Tensor& destination_row_ptr,
+                               const torch::Tensor& source_row_ptr,
+                               const torch::Tensor& source_order,
+                               const torch::Tensor& n_node,
+                               const torch::Tensor& n_local,
+                               const torch::Tensor& atype,
+                               const torch::Tensor& type_embedding,
+                               const torch::Tensor& contact_radius,
+                               double f_inner,
+                               double f_outer) {
+  TORCH_CHECK(edge_vec.is_cuda(), op, ": edge_vec must be a CUDA tensor");
+  source = deepmd::canonical_index_view(source);
+  TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), op,
+              ": source and edge_vec must share the edge axis");
+  const long node_count = atype.size(0);
+  const long edge_count = edge_vec.size(0);
+  TORCH_CHECK(destination_row_ptr.numel() == node_count + 1, op,
+              ": atype and destination_row_ptr describe different node "
+              "counts");
+  const torch::Device device = edge_vec.device();
+  const auto on_device = [&](const torch::Tensor& tensor) {
+    return tensor.is_cuda() && tensor.device() == device &&
+           tensor.is_contiguous();
+  };
+  TORCH_CHECK(on_device(n_node) && n_node.dim() == 1 &&
+                  n_node.scalar_type() == torch::kInt64 && on_device(n_local) &&
+                  n_local.dim() == 1 &&
+                  n_local.scalar_type() == torch::kInt64 &&
+                  n_local.numel() == n_node.numel() && n_node.numel() > 0,
+              op,
+              ": n_node and n_local must be contiguous int64 tensors on the "
+              "device of edge_vec with one shared shape (nf,), nf >= 1");
+  TORCH_CHECK(on_device(atype) && atype.dim() == 1 &&
+                  atype.scalar_type() == torch::kInt64,
+              op,
+              ": atype must be a contiguous int64 tensor of shape (N,) on the "
+              "device of edge_vec");
+  TORCH_CHECK(
+      f_outer <= f_inner || (on_device(contact_radius) &&
+                             contact_radius.scalar_type() == torch::kFloat32 &&
+                             contact_radius.dim() == 1 &&
+                             contact_radius.size(0) == type_embedding.size(0)),
+      op,
+      ": a windowed descriptor needs contiguous fp32 contact radii of shape "
+      "(type_count,) on the device of edge_vec");
+  const bool atom_virial = source_row_ptr.numel() == node_count + 1 &&
+                           source_order.numel() == edge_count;
+  TORCH_CHECK(
+      atom_virial || (source_row_ptr.numel() == 2 && source_order.numel() == 2),
+      op, ": the source CSR must hold ", node_count + 1, " row pointers and ",
+      edge_count,
+      " positions, or two entries each to decline the per-atom virial; got ",
+      source_row_ptr.numel(), " and ", source_order.numel());
+  if (atom_virial) {
+    TORCH_CHECK((source.scalar_type() == torch::kUInt32 ||
+                 source.scalar_type() == torch::kInt64) &&
+                    on_device(source_row_ptr) &&
+                    source_row_ptr.scalar_type() == torch::kInt64 &&
+                    on_device(source_order) &&
+                    source_order.scalar_type() == torch::kInt64,
+                op,
+                ": the per-atom virial needs uint32 or int64 sources and a "
+                "source CSR of contiguous int64 row pointers and positions on "
+                "the device of edge_vec");
+  }
+  return atom_virial;
+}
+
+// A fixed-point force word cannot hold a non-finite value, so a frame of a
+// batch whose energy or virial is not finite reports every force of its
+// nodes, and every magnetic force unless ``force_mag`` is null, as not a
+// number (see energy_force_epilogue_kernel). Nodes past the last frame
+// belong to none and keep their forces.
+void mark_broken_frames(const torch::Tensor& energy,
+                        const torch::Tensor& virial,
+                        const torch::Tensor& frame_row_ptr,
+                        long node_count,
+                        torch::Tensor& force,
+                        torch::Tensor* force_mag) {
+  const auto broken =
+      torch::cat({torch::logical_not(torch::isfinite(energy).all(1).logical_and(
+                      torch::isfinite(virial).flatten(1).all(1))),
+                  torch::zeros({1}, energy.options().dtype(torch::kBool))});
+  const auto frame =
+      torch::searchsorted(frame_row_ptr,
+                          torch::arange(node_count, frame_row_ptr.options()),
+                          /*out_int32=*/false, /*right=*/true) -
+      1;
+  const auto node_broken = broken.index_select(0, frame).unsqueeze(1);
+  force.masked_fill_(node_broken, std::numeric_limits::quiet_NaN());
+  if (force_mag != nullptr) {
+    force_mag->masked_fill_(node_broken,
+                            std::numeric_limits::quiet_NaN());
+  }
+}
+
+// A fitting network of the fused energy-force operators -- the energy fitting
+// or the charge network of a long-range module: its layer plan, weights,
+// head, per-type bias and activation.
+struct FittingNetwork {
+  const FittingLayerPlan& plan;
+  const std::vector& ws;
+  const std::vector& bs;
+  const std::vector& resnets;
+  const torch::Tensor& w_head;
+  const torch::Tensor& b_head;
+  const torch::Tensor& bias_atom_e;
+  int64_t act;
+};
+
+// The fitting routes of the runs of ``run`` nodes into which a fused operator
+// splits its node axis, and the buffers they need. The full runs and the tail
+// run each take the route of their own size; an absent table, like an empty
+// one, leaves both to the rules. A route of the fused network needs the
+// transposed weights, and its kernel writes the pre-activations of a run and
+// the sheet through which the CTAs of a cluster exchange their columns into
+// the pre-activation buffer; the per-layer kernels need the activation
+// ping-pong, sized for the largest run that takes them, and the
+// pre-activations of a run between their forward and their backward. With
+// ``deferred``, every run of the per-layer kernels keeps its pre-activations
+// at its own place on the node axis, so that its backward may follow the
+// forward of other runs and other work of the operator; the buffer then holds
+// every node, as the descriptor of the long-range operator does.
+class FittingRuns {
+ public:
+  FittingRuns(cudaStream_t stream,
+              const FittingNetwork& fitting,
+              long run,
+              long node_count,
+              std::optional> table,
+              const torch::TensorOptions& f32,
+              bool deferred = false)
+      : fitting_(fitting), run_(run) {
+    const std::vector routes =
+        std::move(table).value_or(std::vector{});
+    const long tail = node_count % run;
+    const long full_count = node_count / run;
+    full_route_ = fitting_route(fitting.plan, fitting.ws, run, routes);
+    tail_route_ = tail > 0
+                      ? fitting_route(fitting.plan, fitting.ws, tail, routes)
+                      : full_route_;
+    const long per_layer_nodes =
+        std::max(full_count > 0 && full_route_ == kFittingPerLayer ? run : 0L,
+                 tail > 0 && tail_route_ == kFittingPerLayer ? tail : 0L);
+    if ((full_count > 0 && full_route_ != kFittingPerLayer) ||
+        (tail > 0 && tail_route_ != kFittingPerLayer)) {
+      transposed_ =
+          torch::empty({fitting_fused_transposed_floats(fitting.ws)}, f32);
+      fitting_fused_transpose(stream, fitting.plan, fitting.ws,
+                              transposed_.data_ptr());
+    }
+    if (per_layer_nodes > 0) {
+      const int slots = fitting.plan.n_layer > 1 ? 2 : 1;
+      scratch_ =
+          torch::empty({slots, per_layer_nodes, fitting.plan.width_max}, f32);
+      slot_[0] = scratch_[0].data_ptr();
+      slot_[1] = slots > 1 ? scratch_[1].data_ptr() : slot_[0];
+    }
+    // The kept pre-activations of the run from node ``begin`` start at
+    // ``begin * saved_width``, a 16-byte boundary, since every hidden width is
+    // a multiple of four (check_weights of graph_fitting.cu).
+    kept_ = deferred && per_layer_nodes > 0;
+    saved_ = torch::empty(
+        {(kept_ ? node_count : run) * fitting.plan.saved_width()}, f32);
+  }
+
+  // The outputs ``output`` of the ``count`` nodes of the run from node
+  // ``begin``, whose descriptor rows are ``x``, and the cotangent ``d_x`` of
+  // those rows from the output cotangent ``seed``; ``d_x`` may take the place
+  // of ``x``.
+  void output_and_cotangent(cudaStream_t stream,
+                            long begin,
+                            long count,
+                            const float* x,
+                            float* d_x,
+                            long width,
+                            const long* node_type,
+                            const double* seed,
+                            double* output) const {
+    forward(stream, begin, count, x, d_x, width, node_type, seed, output);
+    backward(stream, begin, count, seed, d_x);
+  }
+
+  // The outputs of a run as in output_and_cotangent, and on a route of the
+  // fused network its cotangent as well, which that network forms in the same
+  // pass; a run of the per-layer kernels saves its pre-activations for
+  // backward().
+  void forward(cudaStream_t stream,
+               long begin,
+               long count,
+               const float* x,
+               float* d_x,
+               long width,
+               const long* node_type,
+               const double* seed,
+               double* output) const {
+    const FittingNetwork& fitting = fitting_;
+    const int64_t route = route_of(count);
+    if (route != kFittingPerLayer) {
+      fitting_energy_gradient_range(
+          stream, fitting.plan, x, width, node_type, fitting.ws, fitting.bs,
+          fitting.resnets, fitting.w_head, fitting.b_head, fitting.bias_atom_e,
+          fitting.act, transposed_.data_ptr(), seed, count, saved(begin),
+          output, d_x, route);
+      return;
+    }
+    fitting_forward_range(stream, fitting.plan, x, width, node_type, fitting.ws,
+                          fitting.bs, fitting.resnets, fitting.w_head,
+                          fitting.b_head, fitting.bias_atom_e, fitting.act,
+                          count, saved(begin), slot_, output);
+  }
+
+  // The cotangent ``d_x`` of a run of the per-layer kernels from ``seed``,
+  // after its forward(); nothing for a run of the fused network, whose
+  // forward() has formed it. Without ``deferred`` it follows the forward of
+  // its own run.
+  void backward(cudaStream_t stream,
+                long begin,
+                long count,
+                const double* seed,
+                float* d_x) const {
+    if (route_of(count) != kFittingPerLayer) {
+      return;
+    }
+    const FittingNetwork& fitting = fitting_;
+    fitting_backward_range(stream, fitting.plan, seed, saved(begin), fitting.ws,
+                           fitting.bs, fitting.resnets, fitting.w_head,
+                           fitting.act, count, slot_[0],
+                           fitting.plan.n_layer > 1 ? slot_[1] : nullptr, d_x);
+  }
+
+ private:
+  int64_t route_of(long count) const {
+    return count == run_ ? full_route_ : tail_route_;
+  }
+
+  float* saved(long begin) const {
+    return saved_.data_ptr() +
+           (kept_ ? begin * fitting_.plan.saved_width() : 0);
+  }
+
+  const FittingNetwork& fitting_;
+  long run_;
+  int64_t full_route_ = 0;
+  int64_t tail_route_ = 0;
+  bool kept_ = false;
+  torch::Tensor transposed_;
+  torch::Tensor scratch_;
+  torch::Tensor saved_;
+  float* slot_[2] = {nullptr, nullptr};
+};
+
+// The state of a fused energy-force evaluation that its epilogue reads.
+struct EpilogueInputs {
+  const torch::Tensor& n_node;
+  const torch::Tensor& frame_row_ptr;
+  const torch::Tensor& force_fixed;
+  const torch::Tensor& raw_energy;
+  const torch::Tensor& seed;
+  torch::Tensor& node_virial;
+  EvaluationWords* evaluation_words;
+};
+
+// The magnetic terms of a native-spin model: the fixed-point totals and the
+// on-site gradient that form the magnetic force.
+struct SpinTerms {
+  const torch::Tensor& fixed;
+  const torch::Tensor& gradient;
+  torch::Tensor& force_mag;
+};
+
+// The inputs of the per-atom virial, which the source CSR attributes to the
+// source nodes of the edges.
+struct SourceVirialTerms {
+  const torch::Tensor& edge_gradient;
+  const torch::Tensor& edge_vec;
+  const torch::Tensor& row_ptr;
+  const torch::Tensor& order;
+};
+
+// The terms of a long-range module: node energies (N, 1) float64, forces at
+// fixed charges (N, 3) and per-node virials (N, 9) float32.
+struct LongRangeTerms {
+  const torch::Tensor& energy;
+  const torch::Tensor& force;
+  const torch::Tensor& atom_virial;
+};
+
+// The node results and the frame totals of a fused energy-force evaluation
+// (energy_force_epilogue_kernel for a single frame, energy_force_nodes_kernel
+// and segment reductions for several), and the per-atom virial
+// ``atom_virial_output``, which is zero without ``source_virial``. ``spin`` is
+// null for a model without spin and ``long_range`` for a model without a
+// long-range module.
+void finish_energy_force(cudaStream_t stream,
+                         long node_count,
+                         int blocks,
+                         const EpilogueInputs& step,
+                         const SpinTerms* spin,
+                         const SourceVirialTerms* source_virial,
+                         const LongRangeTerms* long_range,
+                         torch::Tensor& energy,
+                         torch::Tensor& atom_energy,
+                         torch::Tensor& force,
+                         torch::Tensor& virial,
+                         torch::Tensor& atom_virial_output) {
+  const auto* force_words = reinterpret_cast(
+      step.force_fixed.data_ptr());
+  const auto* spin_words = spin != nullptr
+                               ? reinterpret_cast(
+                                     spin->fixed.data_ptr())
+                               : nullptr;
+  const float* spin_gradient =
+      spin != nullptr ? spin->gradient.data_ptr() : nullptr;
+  double* force_mag =
+      spin != nullptr ? spin->force_mag.data_ptr() : nullptr;
+  const double* extra_energy =
+      long_range != nullptr ? long_range->energy.data_ptr() : nullptr;
+  const float* extra_force =
+      long_range != nullptr ? long_range->force.data_ptr() : nullptr;
+  float* atom_virial_zero =
+      source_virial != nullptr ? nullptr : atom_virial_output.data_ptr();
+  if (step.n_node.size(0) == 1) {
+    auto partial = torch::empty({blocks, kEpilogueTotals},
+                                energy.options().dtype(torch::kFloat64));
+    energy_force_epilogue_kernel<<>>(
+        node_count, step.n_node.data_ptr(), force_words, spin_words,
+        spin_gradient, step.raw_energy.data_ptr(),
+        step.seed.data_ptr(), step.node_virial.data_ptr(),
+        extra_energy, extra_force,
+        long_range != nullptr ? long_range->atom_virial.data_ptr()
+                              : nullptr,
+        force.data_ptr(), force_mag, atom_energy.data_ptr(),
+        atom_virial_zero, partial.data_ptr(), step.evaluation_words,
+        energy.data_ptr(), virial.data_ptr());
+    DPA4C_CHECK_LAUNCH("dpa4c energy-force epilogue");
+  } else {
+    energy_force_nodes_kernel<<>>(
+        node_count, force_words, spin_words, spin_gradient,
+        step.raw_energy.data_ptr(), step.seed.data_ptr(),
+        extra_energy, extra_force, force.data_ptr(), force_mag,
+        atom_energy.data_ptr(), atom_virial_zero);
+    DPA4C_CHECK_LAUNCH("dpa4c energy-force nodes");
+    energy.copy_(frame_scalar_sum(atom_energy, step.n_node));
+    if (long_range != nullptr) {
+      step.node_virial.add_(long_range->atom_virial);
+    }
+    virial.copy_(frame_virial_sum(step.node_virial, step.n_node));
+    mark_broken_frames(energy, virial, step.frame_row_ptr, node_count, force,
+                       spin != nullptr ? &spin->force_mag : nullptr);
+  }
+  if (source_virial != nullptr) {
+    atom_virial_output.copy_(canonical_source_virial(
+        source_virial->edge_gradient, source_virial->edge_vec,
+        source_virial->row_ptr, source_virial->order, node_count));
+    if (long_range != nullptr) {
+      atom_virial_output.add_(
+          long_range->atom_virial.reshape({node_count, 3, 3}));
+    }
+  }
+}
+
+}  // namespace
+
+// Energy, force and virial of one compressed inference step, evaluated over
 // runs of consecutive destination nodes.
 //
 // Destination-sorted CSR gives a run a contiguous span of the edge axis and
 // assigns every edge to exactly one destination, so the runs partition the
 // work rather than splitting any reduction. Folding the descriptor, the
-// fitting and the descriptor backward into one operator lets a run retire its
-// descriptor, its cotangent and its moment state before the next run starts,
-// which leaves only the graph and the edge cotangent at system scale. Nothing
-// is recomputed: the fitting seed is the ownership mask, known up front.
+// fitting, the descriptor backward and the force assembly into one operator
+// lets a run retire its descriptor, its cotangent and its moment state before
+// the next run starts, and unless the per-atom virial is requested no
+// edge-scale cotangent is materialized: the edge backward adds the source
+// half of every force onto the fixed-point totals and reduces the
+// destination half and the virial over its own row.
+// Nothing is recomputed: the fitting seed is the ownership mask, known up
+// front.
+//
+// The operator also owns the node bookkeeping around the network: the
+// ownership mask that selects the energy-contributing nodes, the masked node
+// energies, the frame totals and the conversion of the fixed-point force
+// totals, so that one evaluation of a single frame adds a prologue, an
+// epilogue and the fill that clears their words to its descriptor and fitting
+// kernels.
+//
+// The per-atom virial is attributed to source nodes and is the one output
+// that reads the source CSR. A caller that does not want it passes two-entry
+// placeholders of zeros in its place; the per-atom virial is then returned as
+// zeros and the edge cotangents are never stored. A graph of one node and two
+// edge slots reads such placeholders as a complete source CSR of no edge, with
+// the same result.
+//
+// A force contribution of 2^30 eV/Å or more lies outside the range of the
+// fixed-point force totals, and marks its frame like a non-finite total (see
+// kForceFixedLimit for the scope of the check).
+//
+// Every reduction has a fixed order for a given node count and kernel
+// configuration, so an evaluation is reproducible bit for bit on one device;
+// the configuration follows from the device, so another device may differ in
+// the last bits.
+//
+// The frame counts ``n_node`` sum to at most the node count, and ``n_local``
+// does not exceed them; nodes past their sum belong to no frame and enter no
+// total.
 //
 // The loop lives here rather than in Python because its trip count follows a
 // dynamic node count, which export cannot trace.
-std::tuple
-dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec,
-                                         torch::Tensor source,
-                                         torch::Tensor destination_row_ptr,
-                                         torch::Tensor atype,
-                                         torch::Tensor table,
-                                         torch::Tensor pair_film,
-                                         torch::Tensor pair_mixing,
-                                         torch::Tensor type_embedding,
-                                         torch::Tensor readout_matrices,
-                                         torch::Tensor coupling_meta,
-                                         torch::Tensor coupling_entry,
-                                         torch::Tensor coupling_value,
-                                         torch::Tensor output_mean,
-                                         torch::Tensor output_inv_std,
-                                         torch::Tensor spin,
-                                         torch::Tensor spin_pair,
-                                         torch::Tensor spin_type,
-                                         torch::Tensor contact_radius,
-                                         int64_t lmax,
-                                         double table_stride,
-                                         double table_max,
-                                         double rcut,
-                                         double eps,
-                                         double degree_floor,
-                                         double f_inner,
-                                         double f_outer,
-                                         std::vector ws,
-                                         std::vector bs,
-                                         std::vector resnets,
-                                         torch::Tensor w_head,
-                                         torch::Tensor b_head,
-                                         torch::Tensor bias_atom_e,
-                                         int64_t act,
-                                         torch::Tensor seed,
-                                         int64_t tile,
-                                         torch::Tensor pair_table) {
-  TORCH_CHECK(edge_vec.is_cuda(),
-              "dpa4c_canonical_compress_energy_gradient: edge_vec must be a "
-              "CUDA tensor");
-  TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0),
-              "dpa4c_canonical_compress_energy_gradient: source and edge_vec "
-              "must share the edge axis");
+std::tuple
+dpa4c_canonical_compress_energy_force(
+    torch::Tensor edge_vec,
+    torch::Tensor source,
+    torch::Tensor destination_row_ptr,
+    torch::Tensor source_row_ptr,
+    torch::Tensor source_order,
+    torch::Tensor n_node,
+    torch::Tensor n_local,
+    torch::Tensor atype,
+    torch::Tensor table,
+    torch::Tensor basis_table,
+    torch::Tensor basis_matrix,
+    torch::Tensor pair_film,
+    torch::Tensor pair_mixing,
+    torch::Tensor type_embedding,
+    torch::Tensor readout_matrices,
+    torch::Tensor coupling_meta,
+    torch::Tensor coupling_entry,
+    torch::Tensor coupling_value,
+    torch::Tensor output_mean,
+    torch::Tensor output_inv_std,
+    torch::Tensor spin,
+    torch::Tensor spin_pair,
+    torch::Tensor spin_type,
+    torch::Tensor contact_radius,
+    int64_t lmax,
+    double table_stride,
+    double table_max,
+    double rcut,
+    double eps,
+    double degree_floor,
+    double f_inner,
+    double f_outer,
+    std::vector ws,
+    std::vector bs,
+    std::vector resnets,
+    torch::Tensor w_head,
+    torch::Tensor b_head,
+    torch::Tensor bias_atom_e,
+    int64_t act,
+    int64_t tile,
+    torch::Tensor pair_table,
+    std::optional> fitting_routes) {
+  const bool atom_virial = check_energy_force_inputs(
+      "dpa4c_canonical_compress_energy_force", edge_vec, source,
+      destination_row_ptr, source_row_ptr, source_order, n_node, n_local, atype,
+      type_embedding, contact_radius, f_inner, f_outer);
   const c10::cuda::CUDAGuard device_guard(edge_vec.device());
   const long node_count = atype.size(0);
-  TORCH_CHECK(destination_row_ptr.numel() == node_count + 1,
-              "dpa4c_canonical_compress_energy_gradient: atype and "
-              "destination_row_ptr describe different node counts");
   const int channels = static_cast(type_embedding.size(1));
   const bool has_spin = spin.dim() == 2;
   const Dimensions widths =
       profile_dimensions(channels, static_cast(lmax), has_spin);
-  auto f32 = edge_vec.options().dtype(torch::kFloat32);
+  const long frame_count = n_node.size(0);
+  const auto f32 = edge_vec.options().dtype(torch::kFloat32);
+  const auto f64 = edge_vec.options().dtype(torch::kFloat64);
+  const auto i64 = edge_vec.options().dtype(torch::kInt64);
   auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous();
-  auto energy =
-      torch::empty({node_count, 1}, edge_vec.options().dtype(torch::kFloat64));
-  // With no destination rows every edge slot is padding, and the loop that
-  // would clear it never runs.
-  auto edge_gradient = node_count == 0 ? torch::zeros_like(edge_vec_float)
-                                       : torch::empty_like(edge_vec_float);
-  // Every absent output receives its own allocation, so that no two of the
-  // four unannotated results share storage.
-  const auto absent = [&f32] { return torch::empty({0}, f32); };
+  const auto stream = at::cuda::getCurrentCUDAStream();
+
+  // Every output receives its own allocation, so that no two of the seven
+  // unannotated results share storage.
+  auto energy = torch::empty({frame_count, 1}, f64);
+  auto atom_energy = torch::empty({node_count, 1}, f64);
+  auto force = torch::empty({node_count, 3}, f64);
+  auto virial = torch::empty({frame_count, 3, 3}, f64);
+  auto atom_virial_output = torch::empty({node_count, 3, 3}, f32);
+  auto force_mag =
+      has_spin ? torch::empty({node_count, 3}, f64) : torch::empty({0}, f32);
+  auto mask =
+      torch::empty({node_count}, edge_vec.options().dtype(torch::kInt32));
+
+  auto raw_energy = torch::empty({node_count}, f64);
+  auto seed = torch::empty({node_count}, f64);
+  auto node_type = torch::empty({node_count}, i64);
+  auto force_fixed = torch::empty({node_count, kForceWords}, i64);
+  auto node_virial = torch::empty({node_count, 9}, f32);
   // The on-site magnetic gradient is node local, so a run writes only its own
-  // rows. The neighbour part belongs to source nodes that other runs own, so
-  // it is materialized over the whole edge axis and reduced once afterwards,
-  // exactly like the conservative edge cotangent.
-  auto spin_gradient = has_spin ? torch::empty({node_count, 3}, f32) : absent();
-  auto edge_spin_gradient =
-      has_spin ? (node_count == 0 ? torch::zeros_like(edge_vec_float)
-                                  : torch::empty_like(edge_vec_float))
-               : absent();
+  // rows. The neighbour part belongs to source nodes that other runs own and
+  // accumulates in fixed point like the force.
+  auto spin_gradient =
+      has_spin ? torch::empty({node_count, 3}, f32) : torch::Tensor();
+  auto spin_fixed =
+      has_spin ? torch::empty({node_count, kForceWords}, i64) : torch::Tensor();
+  auto edge_gradient =
+      atom_virial ? torch::empty_like(edge_vec_float) : torch::Tensor();
+  const int blocks = epilogue_blocks(node_count);
+  auto words =
+      torch::zeros({static_cast(sizeof(EvaluationWords) / sizeof(int))},
+                   edge_vec.options().dtype(torch::kInt32));
+  auto* evaluation_words =
+      reinterpret_cast(words.data_ptr());
   if (node_count == 0) {
-    return {energy, edge_gradient.to(edge_vec.scalar_type()), spin_gradient,
-            edge_spin_gradient};
+    energy.zero_();
+    virial.zero_();
+    return {energy,    atom_energy, force, virial, atom_virial_output,
+            force_mag, mask};
   }
-  auto seed_c = seed.contiguous();
-  TORCH_CHECK(
-      seed_c.numel() == node_count && seed_c.scalar_type() == torch::kFloat64,
-      "dpa4c_canonical_compress_energy_gradient: seed must be fp64 "
-      "with one entry per node");
+  const torch::Tensor frame_row_ptr =
+      frame_count > 1 ? frame_row_pointer(n_node) : torch::Tensor();
+  energy_force_prologue_kernel<<>>(
+      node_count, frame_count, n_local.data_ptr(),
+      frame_count > 1 ? frame_row_ptr.data_ptr() : nullptr,
+      atype.data_ptr(), node_type.data_ptr(),
+      seed.data_ptr(), mask.data_ptr(), nullptr,
+      reinterpret_cast(force_fixed.data_ptr()),
+      has_spin ? reinterpret_cast(
+                     spin_fixed.data_ptr())
+               : nullptr,
+      f_outer > f_inner ? contact_radius.data_ptr() : nullptr,
+      evaluation_words);
+  DPA4C_CHECK_LAUNCH("dpa4c energy-force prologue");
 
   const FittingLayerPlan plan = fitting_layer_plan(ws);
   TORCH_CHECK(plan.n_layer > 0 && ws[0].size(0) == widths.output_width,
-              "dpa4c_canonical_compress_energy_gradient: first fitting weight "
+              "dpa4c_canonical_compress_energy_force: first fitting weight "
               "does not match the descriptor width");
   const long run = tile > 0
                        ? std::max(1, std::min(tile, node_count))
                        : node_count;
-  const int slots = plan.n_layer > 1 ? 2 : 1;
   auto descriptor = torch::empty({run, widths.output_width}, f32);
   auto state = torch::empty({run, widths.moment_width + 2}, f32);
-  auto saved = torch::empty({run * plan.saved_width()}, f32);
-  auto scratch = torch::empty({slots, run, plan.width_max}, f32);
-  float* slot[2] = {
-      scratch[0].data_ptr(),
-      slots > 1 ? scratch[1].data_ptr() : scratch[0].data_ptr()};
+  const FittingNetwork fitting{plan,   ws,     bs,          resnets,
+                               w_head, b_head, bias_atom_e, act};
+  const FittingRuns runs(stream, fitting, run, node_count,
+                         std::move(fitting_routes), f32);
 
   auto empty_index = torch::empty({0}, source.options());
   auto empty_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool));
-  auto stream = at::cuda::getCurrentCUDAStream();
   for (long begin = 0; begin < node_count; begin += run) {
     const long count = std::min(run, node_count - begin);
     const Payload payload{
@@ -943,8 +1952,10 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec,
         empty_mask,
         empty_index,
         destination_row_ptr.slice(0, begin, begin + count + 1),
-        atype,
+        node_type,
         table,
+        basis_table,
+        basis_matrix,
         pair_film,
         pair_mixing,
         type_embedding,
@@ -970,29 +1981,35 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec,
         begin};
     Arguments arguments =
         build_arguments(payload, edge_vec_float, channels, widths);
+    arguments.system_node_count = node_count;
     arguments.descriptor = descriptor.data_ptr();
     arguments.state_out = state.data_ptr();
+    arguments.type_mismatch = &evaluation_words->node_types.type_mismatch;
+    arguments.largest_radius = &evaluation_words->node_types.largest_radius;
+    arguments.type_present =
+        present_types(arguments, &evaluation_words->node_types);
     dispatch(channels, false, arguments, stream);
 
-    fitting_forward_range(stream, plan, descriptor.data_ptr(),
-                          widths.output_width, atype.data_ptr() + begin,
-                          ws, bs, resnets, w_head, b_head, bias_atom_e, act,
-                          count, saved.data_ptr(), slot,
-                          energy.data_ptr() + begin);
     // The cotangent replaces the descriptor, which the run no longer needs.
-    fitting_backward_range(
-        stream, plan, seed_c.data_ptr() + begin,
-        saved.data_ptr(), ws, bs, resnets, w_head, act, count, slot[0],
-        plan.n_layer > 1 ? slot[1] : nullptr, descriptor.data_ptr());
+    runs.output_and_cotangent(
+        stream, begin, count, descriptor.data_ptr(),
+        descriptor.data_ptr(), widths.output_width,
+        node_type.data_ptr() + begin, seed.data_ptr() + begin,
+        raw_energy.data_ptr() + begin);
 
     arguments.descriptor_gradient = descriptor.data_ptr();
     arguments.state = state.data_ptr();
     arguments.moment_gradient = state.data_ptr();
-    arguments.edge_gradient = edge_gradient.data_ptr();
+    arguments.force =
+        reinterpret_cast(force_fixed.data_ptr());
+    arguments.node_virial = node_virial.data_ptr() + begin * 9;
+    arguments.edge_gradient =
+        atom_virial ? edge_gradient.data_ptr() : nullptr;
     if (has_spin) {
       // Node-indexed like the energy, so the run addresses its own slice.
       arguments.spin_gradient = spin_gradient.data_ptr() + begin * 3;
-      arguments.edge_spin_gradient = edge_spin_gradient.data_ptr();
+      arguments.spin_force =
+          reinterpret_cast(spin_fixed.data_ptr());
     }
     // Only the final run reaches the reserved edge slots; its row pointer ends
     // at the last physical edge, which is exactly where the padding begins.
@@ -1001,13 +2018,345 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec,
       // The fitting forward has assigned the energy rows of this run, and the
       // seed is the cotangent the fitting backward just consumed.
       attach_pair_term(arguments, pair_table, edge_vec.device(),
-                       seed_c.data_ptr() + begin,
-                       energy.data_ptr() + begin);
+                       seed.data_ptr() + begin,
+                       raw_energy.data_ptr() + begin);
     }
     dispatch(channels, true, arguments, stream);
   }
-  return {energy, edge_gradient.to(edge_vec.scalar_type()), spin_gradient,
-          edge_spin_gradient};
+
+  const EpilogueInputs step{n_node, frame_row_ptr, force_fixed,     raw_energy,
+                            seed,   node_virial,   evaluation_words};
+  const SpinTerms spin_terms{spin_fixed, spin_gradient, force_mag};
+  const SourceVirialTerms source_terms{edge_gradient, edge_vec_float,
+                                       source_row_ptr, source_order};
+  finish_energy_force(stream, node_count, blocks, step,
+                      has_spin ? &spin_terms : nullptr,
+                      atom_virial ? &source_terms : nullptr, nullptr, energy,
+                      atom_energy, force, virial, atom_virial_output);
+  return {energy,    atom_energy, force, virial, atom_virial_output,
+          force_mag, mask};
+}
+
+// Energy, force and virial of one compressed inference step of a model with a
+// latent-charge long-range module.
+//
+// The module reads charges that its network predicts from the descriptor of
+// every node, and the energy of the charges depends on all of them, so the
+// descriptor backward cannot start before every charge is known. The
+// operator evaluates the descriptor of every node with its moment state, and
+// the charges from the descriptor; then the long range of the charges at
+// fixed charges with its derivative with respect to them
+// (deepmd::latent_charge_sum); then the energy fitting with its cotangent, to
+// which the Jacobian row of every charge with respect to the descriptor of
+// its node, the cotangent of a unit seed, scaled by that derivative adds the
+// cotangent of the charge network; and one descriptor backward from the sum.
+// As in dpa4c_canonical_compress_energy_force, the two networks evaluate the
+// node axis in runs of ``tile`` nodes, and each takes the fused network or
+// the per-layer kernels by a fitting-route table of its own
+// (``fitting_routes`` and ``charge_fitting_routes``). A run of the fused
+// charge network forms its Jacobian rows in the pass of its charges, and a
+// run of the per-layer kernels after the energy fitting, from the
+// pre-activations that it keeps for every node; on a V100 at 512 atoms that
+// order shortens the LAMMPS step by 3.3% against forming them before the
+// long range. The Jacobian rows take the width of the descriptor for every
+// node. The edge backward adds every force onto the
+// fixed-point totals as in dpa4c_canonical_compress_energy_force, and the
+// epilogue adds the energies, the forces at fixed charges and the per-node
+// virials of the long range.
+//
+// The charge network has one channel and the layout of an eligible energy
+// fitting; its per-type bias is ``charge_bias``. ``coord``, ``total_charge``,
+// ``box``, ``coefficient``, ``communicator``, ``kernel``, ``channel``,
+// ``prefactor``, ``order`` and ``spacing`` are those of
+// deepmd::latent_charge_sum, whose charged nodes are the owned real nodes of
+// the frames. The per-atom virial and the source CSR follow
+// dpa4c_canonical_compress_energy_force; the descriptor carries no spin
+// condition and the model no analytical pair potential.
+std::tuple
+dpa4c_canonical_compress_long_range_energy_force(
+    torch::Tensor edge_vec,
+    torch::Tensor source,
+    torch::Tensor destination_row_ptr,
+    torch::Tensor source_row_ptr,
+    torch::Tensor source_order,
+    torch::Tensor n_node,
+    torch::Tensor n_local,
+    torch::Tensor atype,
+    torch::Tensor table,
+    torch::Tensor basis_table,
+    torch::Tensor basis_matrix,
+    torch::Tensor pair_film,
+    torch::Tensor pair_mixing,
+    torch::Tensor type_embedding,
+    torch::Tensor readout_matrices,
+    torch::Tensor coupling_meta,
+    torch::Tensor coupling_entry,
+    torch::Tensor coupling_value,
+    torch::Tensor output_mean,
+    torch::Tensor output_inv_std,
+    torch::Tensor spin,
+    torch::Tensor spin_pair,
+    torch::Tensor spin_type,
+    torch::Tensor contact_radius,
+    int64_t lmax,
+    double table_stride,
+    double table_max,
+    double rcut,
+    double eps,
+    double degree_floor,
+    double f_inner,
+    double f_outer,
+    std::vector ws,
+    std::vector bs,
+    std::vector resnets,
+    torch::Tensor w_head,
+    torch::Tensor b_head,
+    torch::Tensor bias_atom_e,
+    int64_t act,
+    int64_t tile,
+    std::vector charge_ws,
+    std::vector charge_bs,
+    std::vector charge_resnets,
+    torch::Tensor charge_w_head,
+    torch::Tensor charge_b_head,
+    torch::Tensor charge_bias,
+    int64_t charge_act,
+    torch::Tensor coord,
+    torch::Tensor total_charge,
+    std::optional box,
+    torch::Tensor coefficient,
+    torch::Tensor communicator,
+    int64_t kernel,
+    int64_t channel,
+    double prefactor,
+    int64_t order,
+    double spacing,
+    std::optional> fitting_routes,
+    std::optional> charge_fitting_routes) {
+  constexpr const char* kOp =
+      "dpa4c_canonical_compress_long_range_energy_force";
+  const bool atom_virial = check_energy_force_inputs(
+      kOp, edge_vec, source, destination_row_ptr, source_row_ptr, source_order,
+      n_node, n_local, atype, type_embedding, contact_radius, f_inner, f_outer);
+  const c10::cuda::CUDAGuard device_guard(edge_vec.device());
+  TORCH_CHECK(spin.dim() != 2, kOp,
+              ": a descriptor with a spin condition is not supported");
+  const long node_count = atype.size(0);
+  TORCH_CHECK(coord.is_cuda() && coord.device() == edge_vec.device() &&
+                  coord.dim() == 2 && coord.size(0) == node_count &&
+                  coord.size(1) == 3 && coord.scalar_type() == torch::kFloat64,
+              kOp,
+              ": coord must be a float64 tensor of shape (N, 3) on the device "
+              "of edge_vec");
+  const int channels = static_cast(type_embedding.size(1));
+  const Dimensions widths = profile_dimensions(channels, static_cast(lmax),
+                                               /*has_spin=*/false);
+  const long frame_count = n_node.size(0);
+  const auto f32 = edge_vec.options().dtype(torch::kFloat32);
+  const auto f64 = edge_vec.options().dtype(torch::kFloat64);
+  const auto i64 = edge_vec.options().dtype(torch::kInt64);
+  auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous();
+  const auto stream = at::cuda::getCurrentCUDAStream();
+
+  // Every output receives its own allocation, so that no two of the six
+  // unannotated results share storage.
+  auto energy = torch::empty({frame_count, 1}, f64);
+  auto atom_energy = torch::empty({node_count, 1}, f64);
+  auto force = torch::empty({node_count, 3}, f64);
+  auto virial = torch::empty({frame_count, 3, 3}, f64);
+  auto atom_virial_output = torch::empty({node_count, 3, 3}, f32);
+  auto mask =
+      torch::empty({node_count}, edge_vec.options().dtype(torch::kInt32));
+  if (node_count == 0) {
+    energy.zero_();
+    virial.zero_();
+    return {energy, atom_energy, force, virial, atom_virial_output, mask};
+  }
+
+  auto raw_energy = torch::empty({node_count}, f64);
+  auto seed = torch::empty({node_count}, f64);
+  auto charged =
+      torch::empty({node_count}, edge_vec.options().dtype(torch::kBool));
+  auto node_type = torch::empty({node_count}, i64);
+  auto force_fixed = torch::empty({node_count, kForceWords}, i64);
+  auto node_virial = torch::empty({node_count, 9}, f32);
+  auto edge_gradient =
+      atom_virial ? torch::empty_like(edge_vec_float) : torch::Tensor();
+  const int blocks = epilogue_blocks(node_count);
+  auto words =
+      torch::zeros({static_cast(sizeof(EvaluationWords) / sizeof(int))},
+                   edge_vec.options().dtype(torch::kInt32));
+  auto* evaluation_words =
+      reinterpret_cast(words.data_ptr());
+  const torch::Tensor frame_row_ptr =
+      frame_count > 1 ? frame_row_pointer(n_node) : torch::Tensor();
+  energy_force_prologue_kernel<<>>(
+      node_count, frame_count, n_local.data_ptr(),
+      frame_count > 1 ? frame_row_ptr.data_ptr() : nullptr,
+      atype.data_ptr(), node_type.data_ptr(),
+      seed.data_ptr(), mask.data_ptr(), charged.data_ptr(),
+      reinterpret_cast(force_fixed.data_ptr()),
+      nullptr, f_outer > f_inner ? contact_radius.data_ptr() : nullptr,
+      evaluation_words);
+  DPA4C_CHECK_LAUNCH("dpa4c long-range energy-force prologue");
+
+  const FittingLayerPlan plan = fitting_layer_plan(ws);
+  const FittingLayerPlan charge_plan = fitting_layer_plan(charge_ws);
+  TORCH_CHECK(plan.n_layer > 0 && ws[0].size(0) == widths.output_width &&
+                  charge_plan.n_layer > 0 &&
+                  charge_ws[0].size(0) == widths.output_width,
+              kOp,
+              ": the first weights of the energy fitting and of the charge "
+              "network must match the descriptor width");
+  TORCH_CHECK(charge_w_head.numel() == charge_ws.back().size(1) &&
+                  charge_bias.dim() == 1 &&
+                  charge_bias.size(0) == bias_atom_e.size(0),
+              kOp,
+              ": the charge network must have a scalar head and one bias per "
+              "type, as the energy fitting has");
+
+  // === Step 1. Descriptor and moment state of every node ===
+  auto descriptor = torch::empty({node_count, widths.output_width}, f32);
+  auto state = torch::empty({node_count, widths.moment_width + 2}, f32);
+  auto empty_index = torch::empty({0}, source.options());
+  auto empty_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool));
+  const Payload payload{source,
+                        empty_mask,
+                        empty_index,
+                        destination_row_ptr,
+                        node_type,
+                        table,
+                        basis_table,
+                        basis_matrix,
+                        pair_film,
+                        pair_mixing,
+                        type_embedding,
+                        readout_matrices,
+                        coupling_meta,
+                        coupling_entry,
+                        coupling_value,
+                        output_mean,
+                        output_inv_std,
+                        spin,
+                        spin_pair,
+                        spin_type,
+                        contact_radius,
+                        true,
+                        lmax,
+                        table_stride,
+                        table_max,
+                        rcut,
+                        eps,
+                        degree_floor,
+                        f_inner,
+                        f_outer,
+                        0};
+  Arguments arguments =
+      build_arguments(payload, edge_vec_float, channels, widths);
+  arguments.system_node_count = node_count;
+  arguments.descriptor = descriptor.data_ptr();
+  arguments.state_out = state.data_ptr();
+  arguments.type_mismatch = &evaluation_words->node_types.type_mismatch;
+  arguments.largest_radius = &evaluation_words->node_types.largest_radius;
+  arguments.type_present =
+      present_types(arguments, &evaluation_words->node_types);
+  dispatch(channels, false, arguments, stream);
+
+  // === Step 2. Charges of the nodes and their descriptor Jacobian ===
+  // The charge of a node is a function of the descriptor of that node alone,
+  // so the cotangent of its descriptor through the charge network is the
+  // derivative of the energy with respect to the charge times the Jacobian
+  // row of the charge, the input cotangent of a unit seed. A run of the fused
+  // network forms the charges and those rows in one pass, before the long
+  // range forms the derivative and the energy fitting overwrites the
+  // descriptor; a run of the per-layer kernels forms its charges here and
+  // its rows from their saved pre-activations after the energy fitting.
+  const long run = tile > 0
+                       ? std::max(1, std::min(tile, node_count))
+                       : node_count;
+  const long width = widths.output_width;
+  const auto charge_bias_double = charge_bias.to(torch::kFloat64).contiguous();
+  const FittingNetwork charge_fitting{
+      charge_plan,   charge_ws,     charge_bs,          charge_resnets,
+      charge_w_head, charge_b_head, charge_bias_double, charge_act};
+  const FittingRuns charge_runs(stream, charge_fitting, run, node_count,
+                                std::move(charge_fitting_routes), f32, true);
+  const auto unit_seed = torch::ones({run}, f64);
+  auto raw_charge = torch::empty({node_count}, f64);
+  auto charge_jacobian = torch::empty({node_count, width}, f32);
+  for (long begin = 0; begin < node_count; begin += run) {
+    const long count = std::min(run, node_count - begin);
+    charge_runs.forward(
+        stream, begin, count, descriptor.data_ptr() + begin * width,
+        charge_jacobian.data_ptr() + begin * width, width,
+        node_type.data_ptr() + begin, unit_seed.data_ptr(),
+        raw_charge.data_ptr() + begin);
+  }
+
+  // === Step 3. Long range of the charges ===
+  // The per-node virials of the charged nodes sum to the virial of each frame,
+  // which the frame totals form with those of the descriptor.
+  const auto long_range = deepmd::long_range::latent_charge_sum(
+      coord, raw_charge.to(torch::kFloat32).unsqueeze(1), charged, n_node,
+      total_charge, box, coefficient, communicator, kernel, channel, prefactor,
+      order, spacing);
+  const torch::Tensor& long_range_energy = std::get<0>(long_range);
+  const torch::Tensor& charge_gradient = std::get<1>(long_range);
+  const torch::Tensor& long_range_force = std::get<2>(long_range);
+  const torch::Tensor& long_range_atom_virial = std::get<3>(long_range);
+
+  // === Step 4. Cotangent of the descriptor ===
+  // The cotangent of the energy fitting replaces the descriptor run by run;
+  // the runs of the per-layer kernels of the charge network then form their
+  // Jacobian rows, and the rows, scaled by the derivative of the energy with
+  // respect to the charges, add the cotangent of the charge network.
+  const FittingNetwork fitting{plan,   ws,     bs,          resnets,
+                               w_head, b_head, bias_atom_e, act};
+  const FittingRuns runs(stream, fitting, run, node_count,
+                         std::move(fitting_routes), f32);
+  for (long begin = 0; begin < node_count; begin += run) {
+    const long count = std::min(run, node_count - begin);
+    float* const run_descriptor = descriptor.data_ptr() + begin * width;
+    runs.output_and_cotangent(
+        stream, begin, count, run_descriptor, run_descriptor, width,
+        node_type.data_ptr() + begin, seed.data_ptr() + begin,
+        raw_energy.data_ptr() + begin);
+  }
+  for (long begin = 0; begin < node_count; begin += run) {
+    const long count = std::min(run, node_count - begin);
+    charge_runs.backward(stream, begin, count, unit_seed.data_ptr(),
+                         charge_jacobian.data_ptr() + begin * width);
+  }
+  descriptor.addcmul_(charge_jacobian, charge_gradient);
+
+  // === Step 5. Descriptor backward and force assembly ===
+  arguments.descriptor_gradient = descriptor.data_ptr();
+  arguments.state = state.data_ptr();
+  arguments.moment_gradient = state.data_ptr();
+  arguments.force =
+      reinterpret_cast(force_fixed.data_ptr());
+  arguments.node_virial = node_virial.data_ptr();
+  arguments.edge_gradient =
+      atom_virial ? edge_gradient.data_ptr() : nullptr;
+  arguments.clear_padding = true;
+  dispatch(channels, true, arguments, stream);
+
+  // === Step 6. Node results and frame totals ===
+  const EpilogueInputs step{n_node, frame_row_ptr, force_fixed,     raw_energy,
+                            seed,   node_virial,   evaluation_words};
+  const SourceVirialTerms source_terms{edge_gradient, edge_vec_float,
+                                       source_row_ptr, source_order};
+  const LongRangeTerms long_range_terms{long_range_energy, long_range_force,
+                                        long_range_atom_virial};
+  finish_energy_force(stream, node_count, blocks, step, nullptr,
+                      atom_virial ? &source_terms : nullptr, &long_range_terms,
+                      energy, atom_energy, force, virial, atom_virial_output);
+  return {energy, atom_energy, force, virial, atom_virial_output, mask};
 }
 
 TORCH_LIBRARY_FRAGMENT(deepmd, library) {
@@ -1020,6 +2369,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) {
                &dpa4c_canonical_compress_backward);
   library.impl("dpa4c_canonical_compress_backward_inplace", torch::kCUDA,
                &dpa4c_canonical_compress_backward_inplace);
-  library.impl("dpa4c_canonical_compress_energy_gradient", torch::kCUDA,
-               &dpa4c_canonical_compress_energy_gradient);
+  library.impl("dpa4c_canonical_compress_energy_force", torch::kCUDA,
+               &dpa4c_canonical_compress_energy_force);
+  library.impl("dpa4c_canonical_compress_long_range_energy_force", torch::kCUDA,
+               &dpa4c_canonical_compress_long_range_energy_force);
 }
diff --git a/source/op/pt/dpa4c/graph_compress.cuh b/source/op/pt/dpa4c/graph_compress.cuh
index 6d211836d2..d7a8d47055 100644
--- a/source/op/pt/dpa4c/graph_compress.cuh
+++ b/source/op/pt/dpa4c/graph_compress.cuh
@@ -2,12 +2,16 @@
 //
 // Device templates for the compressed degree-wise DPA4C descriptor.
 //
-// One warp owns one destination node. The forward scan evaluates the
-// tabulated radial branch and its shared mode profiles, applies the ordered
-// PairFiLM amplitude, and reduces both envelope masses together with every
-// degree-wise moment in a single pass over the destination CSR row. The
-// analytical backward runs one node-readout VJP followed by one edge
-// recomputation scan that reproduces the amplitude instead of storing it.
+// A destination row is scanned by one warp, by several warps of one block in
+// a system too small to fill the device, or in bond-length windows by the
+// warps of a co-scheduled block (graph_compress_kernel.cuh). The forward scan
+// evaluates the tabulated radial branch and its shared mode profiles, applies
+// the ordered PairFiLM amplitude -- per edge, or once per node in the folded
+// scans of a single-type system -- and reduces both envelope masses together
+// with every degree-wise moment over the destination CSR row; the readout
+// follows as a launch of its own. The analytical backward runs one
+// node-readout VJP followed by one edge recomputation scan that reproduces
+// the amplitude instead of storing it.
 //
 // Degrees one and two carry the wide channel blocks and are contracted in
 // closed form. Degrees three and four carry a single channel each, so their
@@ -77,6 +81,42 @@ __device__ __forceinline__ float subwarp_sum(float value, unsigned mask) {
   return value;
 }
 
+/// Butterfly sum over one edge group, which leaves the total on every lane.
+/// The leading lane forms the same expression as `subwarp_sum`.
+template 
+__device__ __forceinline__ float subwarp_allsum(float value, unsigned mask) {
+#pragma unroll
+  for (int offset = Width / 2; offset > 0; offset >>= 1) {
+    value += __shfl_xor_sync(mask, value, offset, Width);
+  }
+  return value;
+}
+
+/// Fixed-point image of one force contribution, as the unsigned word that an
+/// integer atomic adds with two's-complement wraparound.
+__device__ __forceinline__ unsigned long long fixed_point(float value) {
+  return static_cast(
+      __float2ll_rn(value * kForceFixedScale));
+}
+
+/// Set the range word of the node whose fixed-point totals start at ``words``
+/// when ``value`` reaches kForceFixedLimit in magnitude; see kForceFixedLimit
+/// for the scope of the check.
+__device__ __forceinline__ void check_force_range(float value,
+                                                  unsigned long long* words) {
+  if (fabsf(value) >= kForceFixedLimit) {
+    atomicOr(words + kForceRangeWord, 1ull);
+  }
+}
+
+/// Component ``component`` of a Cartesian vector held in registers, chosen by
+/// selects rather than by a runtime index, which would move the vector to
+/// local memory.
+__device__ __forceinline__ float vector_component(const float (&vector)[3],
+                                                  int component) {
+  return component == 0 ? vector[0] : (component == 1 ? vector[1] : vector[2]);
+}
+
 template 
 __device__ __forceinline__ long edge_at_position(
     long position, const index_t* destination_order) {
@@ -354,11 +394,27 @@ struct EdgeGeometry {
   long source;
 };
 
+/// Length from which no bond of the row of a centre with radius
+/// `center_radius` lies inside its bridging window, or positive infinity when
+/// the launch carries no largest radius (see `Arguments::largest_radius`). The
+/// window test of a pair compares the bond with the outer fraction of the sum
+/// of the two radii, and the rounded sum and product are monotonic in the
+/// radius they take, so a bond at least this long fails the test for every
+/// source type.
+__device__ __forceinline__ float window_bound(const Arguments& args,
+                                              float center_radius) {
+  return args.largest_radius != nullptr
+             ? args.f_outer * (center_radius + __ldg(args.largest_radius))
+             : __int_as_float(0x7f800000);
+}
+
 // The length scale of a pair is the sum of the two per-element radii. The
-// centre contributes a block constant, so only the source radius is indexed
+// centre contributes a node constant, so only the source radius is indexed
 // per edge; `EdgeGeometry` keeps the source type, from which the backward
 // recovers the same sum, given that constant, at the cost of one further load
-// and one add.
+// and one add. The radius is read only for a bond shorter than the
+// `window_bound` of the row, so the table position of every longer bond
+// depends on its length alone, as in a model without a window.
 template 
 __device__ __forceinline__ EdgeGeometry
 load_geometry(long edge,
@@ -369,6 +425,7 @@ load_geometry(long edge,
               float f_outer,
               float inverse_f_width,
               float center_radius,
+              float bound,
               const float* contact_radius,
               float eps,
               const float* edge_vec) {
@@ -387,13 +444,16 @@ load_geometry(long edge,
   geometry.seen = geometry.radius;
   geometry.envelope = c3_envelope(geometry.radius, rcut);
   if constexpr (Bridged) {
-    const float contact =
-        center_radius + __ldg(contact_radius + geometry.source_type);
-    if (geometry.radius < f_outer * contact) {
-      const float t = window_parameter(__fdividef(geometry.radius, contact),
-                                       f_inner, inverse_f_width);
-      geometry.seen = clamped_length(t, contact, f_inner, f_outer);
-      geometry.envelope = c3_envelope(geometry.seen, rcut) * bridging_switch(t);
+    if (geometry.radius < bound) {
+      const float contact =
+          center_radius + __ldg(contact_radius + geometry.source_type);
+      if (geometry.radius < f_outer * contact) {
+        const float t = window_parameter(__fdividef(geometry.radius, contact),
+                                         f_inner, inverse_f_width);
+        geometry.seen = clamped_length(t, contact, f_inner, f_outer);
+        geometry.envelope =
+            c3_envelope(geometry.seen, rcut) * bridging_switch(t);
+      }
     }
   }
   return geometry;
@@ -416,22 +476,26 @@ template 
 __device__ __forceinline__ float radial_chain(const EdgeGeometry& geometry,
                                               const Arguments& args,
                                               float center_radius,
+                                              float bound,
                                               float table_gradient,
                                               float envelope_gradient) {
   if constexpr (Bridged) {
-    const float contact =
-        center_radius + __ldg(args.contact_radius + geometry.source_type);
-    if (geometry.radius < args.f_outer * contact) {
-      const float t = window_parameter(__fdividef(geometry.radius, contact),
-                                       args.f_inner, args.inverse_f_width);
-      const float seen = clamped_length(t, contact, args.f_inner, args.f_outer);
-      const float weight = bridging_switch(t);
-      const float slope = fmaf(
-          c3_envelope_derivative(seen, args.rcut) * weight, weight,
-          c3_envelope(seen, args.rcut) *
-              __fdividef(bridging_switch_derivative(t, args.inverse_f_width),
-                         contact));
-      return fmaf(envelope_gradient, slope, table_gradient * weight);
+    if (geometry.radius < bound) {
+      const float contact =
+          center_radius + __ldg(args.contact_radius + geometry.source_type);
+      if (geometry.radius < args.f_outer * contact) {
+        const float t = window_parameter(__fdividef(geometry.radius, contact),
+                                         args.f_inner, args.inverse_f_width);
+        const float seen =
+            clamped_length(t, contact, args.f_inner, args.f_outer);
+        const float weight = bridging_switch(t);
+        const float slope = fmaf(
+            c3_envelope_derivative(seen, args.rcut) * weight, weight,
+            c3_envelope(seen, args.rcut) *
+                __fdividef(bridging_switch_derivative(t, args.inverse_f_width),
+                           contact));
+        return fmaf(envelope_gradient, slope, table_gradient * weight);
+      }
     }
   }
   return fmaf(envelope_gradient,
@@ -826,6 +890,70 @@ __device__ __forceinline__ float load_output_gradient(
          __ldg(inverse_stddev + coordinate);
 }
 
+/// Copy ``Count`` floats from global ``source`` to shared ``target``, dealing
+/// them to ``Width`` lanes from ``lane``. A lane issues all of its loads before
+/// its first store, so the latencies of the loads overlap instead of adding
+/// up along the copy.
+template 
+__device__ __forceinline__ void stage_shared(float* target,
+                                             const float* __restrict__ source,
+                                             int lane) {
+  constexpr int Rounds = (Count + Width - 1) / Width;
+  float value[Rounds];
+#pragma unroll
+  for (int round = 0; round < Rounds; ++round) {
+    const int index = lane + round * Width;
+    value[round] = index < Count ? __ldg(source + index) : 0.0f;
+  }
+#pragma unroll
+  for (int round = 0; round < Rounds; ++round) {
+    const int index = lane + round * Width;
+    if (index < Count) {
+      target[index] = value[round];
+    }
+  }
+}
+
+/// Cache level a prefetch request fills.
+enum class CacheLevel { L1, L2 };
+
+/// Request the 128-byte lines that hold ``count`` floats from ``base`` into
+/// `Level`, dealing the lines to the ``width`` lanes from ``lane``. An empty
+/// range requests nothing.
+template 
+__device__ __forceinline__ void prefetch_lines(const float* base,
+                                               long count,
+                                               int lane,
+                                               int width) {
+  if (count <= 0) {
+    return;
+  }
+  const auto first =
+      reinterpret_cast(base) & ~std::uintptr_t{127};
+  const auto last = reinterpret_cast(base + count - 1);
+  const long lines = static_cast((last - first) / 128) + 1;
+  for (long line = lane; line < lines; line += width) {
+    const std::uintptr_t address = first + line * 128ull;
+    if constexpr (Level == CacheLevel::L1) {
+      asm volatile("prefetch.global.L1 [%0];" ::"l"(address));
+    } else {
+      asm volatile("prefetch.global.L2 [%0];" ::"l"(address));
+    }
+  }
+}
+
+/// Request the bond vectors of the canonical destination row ``[begin, end)``
+/// into L2 when the launch streams them (see `Arguments::prefetch_bonds`).
+__device__ __forceinline__ void prefetch_row_bonds(const Arguments& args,
+                                                   long begin,
+                                                   long end,
+                                                   int lane) {
+  if (args.prefetch_bonds) {
+    prefetch_lines(args.edge_vec + 3 * begin, 3 * (end - begin),
+                                   lane, kWarpSize);
+  }
+}
+
 template 
 __device__ __forceinline__ float readout_weight(const float* matrices,
                                                 int matrix,
diff --git a/source/op/pt/dpa4c/graph_compress_c128.cu b/source/op/pt/dpa4c/graph_compress_c128.cu
deleted file mode 100644
index f64b8c3309..0000000000
--- a/source/op/pt/dpa4c/graph_compress_c128.cu
+++ /dev/null
@@ -1,12 +0,0 @@
-// SPDX-License-Identifier: LGPL-3.0-or-later
-//
-// Compiled specializations of the compressed DPA4C descriptor for a scalar
-// width of 128 channels.
-
-#include "graph_compress_kernel.cuh"
-
-namespace deepmd_dpa4c {
-
-DPA4C_DEFINE_CHANNEL(128)
-
-}  // namespace deepmd_dpa4c
diff --git a/source/op/pt/dpa4c/graph_compress_c16.cu b/source/op/pt/dpa4c/graph_compress_c16.cu
deleted file mode 100644
index 946c788640..0000000000
--- a/source/op/pt/dpa4c/graph_compress_c16.cu
+++ /dev/null
@@ -1,12 +0,0 @@
-// SPDX-License-Identifier: LGPL-3.0-or-later
-//
-// Compiled specializations of the compressed DPA4C descriptor for a scalar
-// width of 16 channels.
-
-#include "graph_compress_kernel.cuh"
-
-namespace deepmd_dpa4c {
-
-DPA4C_DEFINE_CHANNEL(16)
-
-}  // namespace deepmd_dpa4c
diff --git a/source/op/pt/dpa4c/graph_compress_c32.cu b/source/op/pt/dpa4c/graph_compress_c32.cu
deleted file mode 100644
index f2c28beedf..0000000000
--- a/source/op/pt/dpa4c/graph_compress_c32.cu
+++ /dev/null
@@ -1,12 +0,0 @@
-// SPDX-License-Identifier: LGPL-3.0-or-later
-//
-// Compiled specializations of the compressed DPA4C descriptor for a scalar
-// width of 32 channels.
-
-#include "graph_compress_kernel.cuh"
-
-namespace deepmd_dpa4c {
-
-DPA4C_DEFINE_CHANNEL(32)
-
-}  // namespace deepmd_dpa4c
diff --git a/source/op/pt/dpa4c/graph_compress_c64.cu b/source/op/pt/dpa4c/graph_compress_c64.cu
deleted file mode 100644
index 3ec9a32ecd..0000000000
--- a/source/op/pt/dpa4c/graph_compress_c64.cu
+++ /dev/null
@@ -1,12 +0,0 @@
-// SPDX-License-Identifier: LGPL-3.0-or-later
-//
-// Compiled specializations of the compressed DPA4C descriptor for a scalar
-// width of 64 channels.
-
-#include "graph_compress_kernel.cuh"
-
-namespace deepmd_dpa4c {
-
-DPA4C_DEFINE_CHANNEL(64)
-
-}  // namespace deepmd_dpa4c
diff --git a/source/op/pt/dpa4c/graph_compress_c8.cu b/source/op/pt/dpa4c/graph_compress_c8.cu
deleted file mode 100644
index 1bfafef6a7..0000000000
--- a/source/op/pt/dpa4c/graph_compress_c8.cu
+++ /dev/null
@@ -1,12 +0,0 @@
-// SPDX-License-Identifier: LGPL-3.0-or-later
-//
-// Compiled specializations of the compressed DPA4C descriptor for a scalar
-// width of 8 channels.
-
-#include "graph_compress_kernel.cuh"
-
-namespace deepmd_dpa4c {
-
-DPA4C_DEFINE_CHANNEL(8)
-
-}  // namespace deepmd_dpa4c
diff --git a/source/op/pt/dpa4c/graph_compress_cpu.cc b/source/op/pt/dpa4c/graph_compress_cpu.cc
index 658aed0759..a63450901a 100644
--- a/source/op/pt/dpa4c/graph_compress_cpu.cc
+++ b/source/op/pt/dpa4c/graph_compress_cpu.cc
@@ -411,6 +411,9 @@ void run_scan(ScanFunction scan,
                    });
 }
 
+// The CPU scans read the radial table directly. The low-rank table that the
+// CUDA single-type scan reads instead belongs to the shared schema and is
+// accepted unused.
 std::tuple forward(
     torch::Tensor edge_vec,
     torch::Tensor edge_index,
@@ -419,6 +422,8 @@ std::tuple forward(
     torch::Tensor destination_row_ptr,
     torch::Tensor atype,
     torch::Tensor table,
+    torch::Tensor /*basis_table*/,
+    torch::Tensor /*basis_matrix*/,
     torch::Tensor pair_film,
     torch::Tensor pair_mixing,
     torch::Tensor type_embedding,
@@ -478,6 +483,8 @@ std::tuple backward(
     torch::Tensor destination_row_ptr,
     torch::Tensor atype,
     torch::Tensor table,
+    torch::Tensor /*basis_table*/,
+    torch::Tensor /*basis_matrix*/,
     torch::Tensor pair_film,
     torch::Tensor pair_mixing,
     torch::Tensor type_embedding,
diff --git a/source/op/pt/dpa4c/graph_compress_kernel.cuh b/source/op/pt/dpa4c/graph_compress_kernel.cuh
index a7184af891..ab7960b8b1 100644
--- a/source/op/pt/dpa4c/graph_compress_kernel.cuh
+++ b/source/op/pt/dpa4c/graph_compress_kernel.cuh
@@ -2,14 +2,20 @@
 //
 // Kernels of the compressed degree-wise DPA4C descriptor.
 //
-// This header is included by exactly one translation unit per scalar width,
-// which instantiates the angular-degree and topology specializations of that
-// width. Keeping the instantiations separated bounds the compile time of any
-// single translation unit and lets them proceed in parallel.
+// The header is included by the generated shards of shard.cu.in, each of
+// which instantiates both directions of the descriptor for one scalar width
+// and one or more angular degrees. Keeping the instantiations separated
+// bounds the compile time of any single translation unit and lets them
+// proceed in parallel.
 
 #pragma once
 
-#include 
+#include 
+
+#include 
+#include 
+#include 
+#include 
 
 #include "graph_compress.cuh"
 
@@ -73,6 +79,898 @@ __device__ __forceinline__ void decode_spin_tensor_pair(int entry,
   column = 1;
 }
 
+// === Single-type systems ===
+//
+// The PairFiLM amplitude and the mode mixing are affine in the radial
+// features, with coefficients fixed by the ordered type pair of the edge, and
+// every moment is linear in the amplitude. In a system whose atoms share one
+// type every edge carries the same pair, so the coefficients are constants:
+// the scans then carry the bare radial features and the coefficients apply
+// once per node, to the reduced sums in the forward and to the cotangents in
+// the backward. The extra features of that formulation are the mode profiles,
+// at feature index `r < radial_modes`, and the constant feature that the FiLM
+// shift multiplies, at index `radial_modes`; each reaches every harmonic
+// component.
+//
+// The folded scans pay off once the per-channel coefficient work is a sizeable
+// part of an edge, which holds from thirty-two channels on.
+
+template 
+constexpr bool foldable() {
+  return Canonical && !HasSpin && Channels >= 32;
+}
+
+// === Systems of a few types ===
+//
+// The edges of a row belong to at most as many ordered pairs as the system
+// has node types, and every moment is linear in the coefficients of the pair
+// of each edge. The scans of a system of at most eight types therefore
+// accumulate the features that a pair coefficient multiplies separately for
+// the edges of each source type, in the slot of that type, and apply the
+// coefficients of the pair of each slot once per row: to the reduced sums in
+// the forward and to the node cotangents in the backward. Two kinds of
+// feature are slotted:
+//
+// - The mode profiles. The mode mixing adds `sum_r M[p, c, r] q_r` to the
+//   amplitude of channel `c` of an edge of pair `p`. A slot holds one moment
+//   per mode and harmonic component, and lane `l` of an edge group owns mode
+//   `l % kFoldedModes` of slot `l / kFoldedModes + k * (EdgeWidth /
+//   kFoldedModes)` in its slot tile `k` (`slot_owner`, `owned_slot`).
+// - The basis functions of the low-rank table (see `low_rank_basis`). The
+//   channels from `C1` on feed degree zero alone, and the FiLM amplitude
+//   `gamma[p, c] sum_k V[k, c] b_k + beta[p, c]` of such a channel is linear
+//   in the basis functions `b_k` and in the constant feature that the shift
+//   multiplies. In the forward a lane holds the sums of its basis columns for
+//   every slot, and lane `s` of an edge group the sum of the envelopes of the
+//   edges of slot `s`; the backward reads the cotangents of the basis
+//   functions and of the shift of the slot of every edge from shared memory.
+//
+// The channels below `C1` keep their FiLM coefficients on the edges, as the
+// pairwise scans do.
+//
+// Every slot adds state to every lane, sums in registers and, in the
+// backward, rows of cotangents in shared memory, and every edge passes the
+// slots of its lanes in predicated or zero-weighted updates, so a scan pays
+// for the slots it compiles whatever the system: a system of up to four types
+// takes the kernels of four slots, and one of five to eight types those of
+// eight (`type_slots`). The profiles that slot the mode profiles run
+// the kernels of eight slots in both directions, and those that slot the
+// basis functions alone in the forward only: their backward of eight slots is
+// no faster than the pairwise backward, since the shared rows of eight slots
+// per warp need a shared-memory configuration that takes L1 cache from the
+// table reads, unless the kernel gives up resident warps. The forward of such
+// a system reads the low-rank table and its backward the radial table, two
+// interpolants whose knot values and slopes agree to the tolerance of the
+// adoption check of the low-rank table (`build_low_rank_table`), so its forces
+// follow the gradient of its energy to the accuracy of the compression.
+//
+// A spin profile slots its mode profiles alone, since it carries no low-rank
+// table; the spin families weight the moment of every neighbour and stay on
+// the edges. It has no single-type scan, so its slotted scans serve a system
+// of one type as well, in a single slot.
+
+/// Whether the edge kernels of a profile slot the features of a system of a
+/// few types: a spin-free profile whose single-type scans fold and that
+/// carries radial modes or admits a low-rank table, or a canonical spin
+/// profile of thirty-two channels or more that carries radial modes.
+template 
+constexpr bool slottable() {
+  if constexpr (HasSpin) {
+    return HasModes && Canonical && Channels >= 32;
+  } else {
+    return (HasModes || low_rank_basis(Channels) > 0) &&
+           foldable();
+  }
+}
+
+/// Slot tiles of a lane of an edge group of `EdgeWidth` lanes in a scan of
+/// `Slots` slots.
+template 
+constexpr int slot_tiles() {
+  return Slots * kFoldedModes / EdgeWidth;
+}
+
+/// Lane of an edge group of `EdgeWidth` lanes that holds mode `mode` of slot
+/// `slot`, in its slot tile `slot / (EdgeWidth / kFoldedModes)`.
+template 
+__device__ __forceinline__ int slot_owner(int slot, int mode) {
+  return (slot % (EdgeWidth / kFoldedModes)) * kFoldedModes + mode;
+}
+
+/// Slot whose mode `lane % kFoldedModes` lane `lane` of an edge group of
+/// `EdgeWidth` lanes holds in its slot tile `tile`.
+template 
+__device__ __forceinline__ int owned_slot(int lane, int tile) {
+  return lane / kFoldedModes + tile * (EdgeWidth / kFoldedModes);
+}
+
+/// The `Slots` slots of the node types of a system: byte `s` of `types` holds
+/// the type of slot `s`, the types in increasing order, every byte past the
+/// types of the system holds 0xFF, above every type, and `count` is the
+/// number of types.
+template 
+struct TypeSlots {
+  using Word = std::conditional_t<(Slots > 4), unsigned long long, unsigned>;
+  static_assert(Slots <= static_cast(sizeof(Word)),
+                "a byte of the word holds the type of each slot");
+  Word types = ~Word{0};
+  int count = 0;
+};
+static_assert(32 * kTypeWords <= 0xFF,
+              "every node type of the bitmap lies below the 0xFF sentinel");
+
+/// Number of node types of a system, or zero for a launch without their
+/// bitmap (see `Arguments::type_present`).
+__device__ __forceinline__ int node_type_count(const Arguments& args) {
+  if (args.type_present == nullptr) {
+    return 0;
+  }
+  int count = 0;
+#pragma unroll
+  for (int word = 0; word < kTypeWords; ++word) {
+    count += __popc(__ldg(args.type_present + word));
+  }
+  return count;
+}
+
+/// Whether the system of `args` holds `fewest` to `most` node types, for a
+/// launch that carries their bitmap. The test reads the bitmap alone, which
+/// no pairwise or single-type scan reads, so that the kernel that it retires
+/// keeps the code of a kernel without it (see `retire_unless`).
+__device__ __forceinline__ bool holds_types(const Arguments& args,
+                                            int fewest,
+                                            int most) {
+  const int count = node_type_count(args);
+  return count >= fewest && count <= most;
+}
+
+/// End the calling thread unless `serves` holds; `serves` is uniform over the
+/// launch. The exit is a predicated instruction of inline PTX, which the
+/// optimizer treats as opaque, so the scan that follows is optimized as in a
+/// kernel without the test.
+__device__ __forceinline__ void retire_unless(bool serves) {
+  asm volatile(
+      "{\n\t.reg .pred retire;\n\t"
+      "setp.eq.u32 retire, %0, 0;\n\t"
+      "@retire exit;\n\t}" ::"r"(static_cast(serves)));
+}
+
+/// The slots of the node types of a system of at most `Slots` types.
+template 
+__device__ __forceinline__ TypeSlots load_type_slots(
+    const Arguments& args) {
+  using Word = typename TypeSlots::Word;
+  TypeSlots slots;
+#pragma unroll
+  for (int word = 0; word < kTypeWords; ++word) {
+    for (unsigned bits = __ldg(args.type_present + word); bits != 0u;
+         bits &= bits - 1u) {
+      const Word type = 32u * word + __ffs(bits) - 1u;
+      const int shift = 8 * slots.count;
+      slots.types = (slots.types & ~(Word{0xFFu} << shift)) | (type << shift);
+      ++slots.count;
+    }
+  }
+  return slots;
+}
+
+/// Request the source indices of the canonical destination row `[begin, end)`
+/// into L2 when the launch streams its bonds (see `Arguments::prefetch_bonds`),
+/// one 128-byte line per lane of the warp from `lane`. A slotted scan reads the
+/// type of the source of every edge to select its slot, so the index heads the
+/// dependency chain of the slot sums as the bond vector heads that of the
+/// table reads.
+template 
+__device__ __forceinline__ void prefetch_row_sources(const Arguments& args,
+                                                     const index_t* edge_index,
+                                                     long begin,
+                                                     long end,
+                                                     int lane) {
+  if (args.prefetch_bonds) {
+    prefetch_lines(
+        reinterpret_cast(edge_index + begin),
+        (end - begin) * static_cast(sizeof(index_t) / sizeof(float)),
+        lane, kWarpSize);
+  }
+}
+
+/// Slot of `type`, a node type of the system: the number of its types below
+/// it.
+template 
+__device__ __forceinline__ int type_slot(const TypeSlots& slots,
+                                         int type) {
+  int slot = 0;
+#pragma unroll
+  for (int index = 1; index < Slots; ++index) {
+    slot += type >= static_cast((slots.types >> (8 * index)) & 0xFFu);
+  }
+  return slot;
+}
+
+/// Node type of slot `slot` of a system of more than `slot` types.
+template 
+__device__ __forceinline__ int slot_type(const TypeSlots& slots,
+                                         int slot) {
+  return static_cast((slots.types >> (8 * slot)) & 0xFFu);
+}
+
+/// Edge body of a destination row. A row of varying pairs reads the ordered
+/// pair of every edge (`Pairwise`). A single-type row folds the pair
+/// coefficients onto the node and reads the channels of the radial table
+/// (`Direct`), or the low-rank table of a descriptor that carries one
+/// (`LowRank`, see `low_rank_basis`). A row of a system of a few types slots
+/// its mode profiles and reads the radial table (`Slotted4` and `Slotted8`),
+/// or slots its mode profiles and the basis functions of the low-rank table
+/// (`SlottedLowRank4` and `SlottedLowRank8`), in four or eight slots. Each is
+/// an instantiation of one edge body, so that each loop keeps only its own
+/// state live.
+///
+/// An edge kernel is instantiated for one folded scan: a single-type scan,
+/// which it runs on a single-type system and the pairwise scan on any other,
+/// or a slotted scan, which it runs on every system it serves. Whether a
+/// descriptor carries the low-rank table is known on the host, so the kernels
+/// hold the low-rank instantiations only when they run them.
+enum class Scan {
+  Pairwise,
+  Direct,
+  LowRank,
+  Slotted4,
+  SlottedLowRank4,
+  Slotted8,
+  SlottedLowRank8
+};
+
+template 
+using ScanTag = std::integral_constant;
+
+/// Whether a scan of kind `kind` slots the basis functions of the low-rank
+/// table.
+constexpr bool slots_basis(Scan kind) {
+  return kind == Scan::SlottedLowRank4 || kind == Scan::SlottedLowRank8;
+}
+
+/// Whether a scan of kind `kind` slots the features of a system of a few
+/// types.
+constexpr bool slots_types(Scan kind) {
+  return kind == Scan::Slotted4 || kind == Scan::Slotted8 || slots_basis(kind);
+}
+
+/// Whether a scan of kind `kind` reads the low-rank table.
+constexpr bool reads_low_rank(Scan kind) {
+  return kind == Scan::LowRank || slots_basis(kind);
+}
+
+/// Slots of a scan of kind `kind`: eight for the slotted scans of a system
+/// of five to eight node types, four for the others.
+constexpr int type_slots(Scan kind) {
+  return kind == Scan::Slotted8 || kind == Scan::SlottedLowRank8 ? 8 : 4;
+}
+
+/// Fewest node types of a system that the slotted scans of kind `kind` serve:
+/// one more than the four slots of the other slotted scans for the scans of
+/// eight slots, and for those of four two, or one for a spin profile, which
+/// has no single-type scan.
+template 
+constexpr int fewest_types(Scan kind) {
+  return type_slots(kind) > 4 ? 5 : (HasSpin ? 1 : 2);
+}
+
+/// Table columns a scan of `Kind` reads before its mode profiles or extra
+/// features: the channels, or the degree-one channels followed by the basis
+/// functions.
+template 
+constexpr int fold_columns() {
+  return reads_low_rank(Kind) ? P::C1 + P::Basis : P::C0;
+}
+
+/// Whether the edge scans of direction `Backward` of a descriptor whose
+/// single-type scan is `Fold` launch the slotted kernels of `Slots` slots
+/// (see `type_slots`) beside their other kernel: for a profile that slots,
+/// those of four slots when the descriptor carries the mode profiles or the
+/// low-rank table, and those of eight slots as well but in the backward of a
+/// descriptor without mode profiles.
+template 
+constexpr bool runs_slotted() {
+  return slottable() &&
+         (HasModes || Fold == Scan::LowRank) &&
+         (Slots <= 4 || HasModes || !Backward);
+}
+
+/// Whether the edge kernel of folded scan `Fold` of direction `Backward`
+/// retires on the systems that another kernel of its launch serves: every
+/// slotted kernel, and the other kernel of a launch with slotted kernels.
+template 
+constexpr bool shares_launch() {
+  return slots_types(Fold) || runs_slotted();
+}
+
+/// Whether the edge kernel of folded scan `Fold` of direction `Backward`
+/// serves the system of `args`: a slotted kernel the systems of
+/// `fewest_types` to `type_slots` node types, and the other kernel of its
+/// launch every system that no slotted kernel of the launch serves.
+template 
+__device__ __forceinline__ bool serves_system(const Arguments& args) {
+  if constexpr (slots_types(Fold)) {
+    return holds_types(args, fewest_types(Fold), type_slots(Fold));
+  } else {
+    constexpr int Most = runs_slotted()
+                             ? type_slots(Scan::Slotted8)
+                             : type_slots(Scan::Slotted4);
+    return !holds_types(args, fewest_types(Scan::Slotted4), Most);
+  }
+}
+
+/// Whether the edge kernels of a profile run the low-rank single-type scan
+/// for a descriptor that carries the low-rank table.
+template 
+constexpr bool folds_low_rank() {
+  return Profile::Basis > 0 &&
+         foldable();
+}
+
+/// Check the folded scan `Fold` of an edge kernel of direction `Backward`, and
+/// retire the blocks of the kernel unless it serves the system of its launch:
+/// the slotted kernels serve the systems of a few node types and the other
+/// kernel of their launch every other system (see `retire_unless`).
+template 
+__device__ __forceinline__ void enter_edge_scan(const Arguments& args) {
+  static_assert(Fold != Scan::Pairwise &&
+                    (!reads_low_rank(Fold) ||
+                     folds_low_rank()) &&
+                    (!slots_types(Fold) || slots_basis(Fold) || HasModes),
+                "the folded scan of a kernel is a single-type or a slotted "
+                "scan, reads the low-rank table only where it folds, and "
+                "slots the radial table for the mode profiles");
+  if constexpr (shares_launch()) {
+    retire_unless(
+        serves_system(
+            args));
+  }
+}
+
+/// Run `launch` with the single-type scan of the descriptor in `args`.
+template 
+void dispatch_fold(const Arguments& args, Launch&& launch) {
+  if constexpr (folds_low_rank()) {
+    if (args.basis_table != nullptr) {
+      launch(ScanTag{});
+      return;
+    }
+  }
+  launch(ScanTag{});
+}
+
+/// Rows of a co-scheduled block of a spin edge scan: twenty warps, whose
+/// block of 640 threads keeps 96 registers per thread for the spin payload,
+/// which spills at the eighty registers of a block of `Lockstep` rows.
+constexpr int kSpinLockstepRows = 20;
+
+/// Co-scheduled width of a canonical edge scan, in warps per block: that of
+/// `Lockstep`, in blocks of `kSpinLockstepRows` rows for a spin profile.
+template 
+constexpr int coscheduled_warps() {
+  constexpr int Warps = Lockstep::Warps;
+  if constexpr (!Canonical || Warps == 1) {
+    return 1;
+  } else {
+    return HasSpin ? kSpinLockstepRows : Warps;
+  }
+}
+
+/// Rows of a co-scheduled block of a slotted edge scan of 128 registers per
+/// thread: sixteen warps, a block of 512 threads.
+constexpr int kSlotLockstepRows = 16;
+
+/// Co-scheduled width of a slotted edge scan of `Slots` slots in a system
+/// that fills the device with its blocks, whose rows read the same windows of
+/// their wide table at the same time. The spin-free scans of 128 channels
+/// take the forward blocks of `Lockstep` and backward blocks of
+/// `kSlotLockstepRows` rows, whose slots add their cotangents to the state of
+/// every lane, and the spin scans of 128 channels blocks of
+/// `kSpinLockstepRows` rows in both directions; the forward of eight slots
+/// of 128 channels, whose slot sums exceed the registers of those blocks,
+/// takes blocks of `kSlotLockstepRows` rows with or without spin. The spin
+/// backward of degree three and four of sixty-four channels, which
+/// co-schedules as its pairwise backward does (see `Lockstep`), takes blocks
+/// of `kSlotLockstepRows` rows. The other scans hold the slots in the
+/// registers of `slot_resident_warps` resident warps, in blocks of independent
+/// rows.
+template 
+constexpr int slot_coscheduled_warps() {
+  if constexpr (Channels >= 128 && Slots > 4 && !Backward) {
+    return kSlotLockstepRows;
+  } else if constexpr (HasSpin) {
+    if constexpr (Channels >= 128) {
+      return kSpinLockstepRows;
+    } else {
+      return Lockstep::Warps > 1 ? kSlotLockstepRows
+                                                           : 1;
+    }
+  } else if constexpr (Channels < 128) {
+    return 1;
+  } else {
+    return Backward ? kSlotLockstepRows
+                    : Lockstep::Warps;
+  }
+}
+
+/// Rows of a block of independent rows of a slotted or a spin edge scan,
+/// each a warp of its own. A kernel that does not serve the system retires
+/// its blocks, whose count sets the time it takes, and every block holds the
+/// shared memory that the driver reserves per block besides its own: four
+/// rows a block retire in a quarter of the time of one, and keep the shared
+/// memory of a multiprocessor of resident spin scans within the configuration
+/// that leaves the most L1, while the rows stay as independent as in blocks
+/// of one row.
+constexpr int kSlotRows = 4;
+
+/// Whether the `Warps` rows of a block walk in the lockstep windows of the
+/// co-scheduled blocks (see `Lockstep`), rather than each on its own: every
+/// block of several rows but the blocks of `kSlotRows` independent rows.
+template 
+constexpr bool windowed() {
+  return Warps > 1 && Warps != kSlotRows;
+}
+static_assert(kSlotRows != Lockstep<64, 3, true>::Warps &&
+                  kSlotRows != Lockstep<128, 2, false>::Warps &&
+                  kSlotRows != kSlotLockstepRows &&
+                  kSlotRows != kSpinLockstepRows,
+              "a block of independent slotted rows is not a co-scheduled one");
+
+/// Launch one edge scan with `Warps` destination rows per block, each row
+/// shared by `Split` warps.
+template 
+void launch_edge_scan(Kernel kernel,
+                      const Arguments& args,
+                      cudaStream_t stream) {
+  kernel<<((args.node_count + Warps - 1) / Warps),
+           Warps * Split * kWarpSize, 0, stream>>>(args);
+}
+
+/// Throw the error of a failed CUDA query that a launch choice derives from,
+/// after clearing it, so that no later launch check reports it again.
+inline void check_query(cudaError_t status, const char* query) {
+  if (status != cudaSuccess) {
+    cudaGetLastError();
+    TORCH_CHECK(false, "dpa4c: ", query,
+                " failed: ", cudaGetErrorString(status));
+  }
+}
+
+/// Per-device memo of a quantity that follows from the current device and the
+/// compiled kernels alone. Each device answers once, and every function that
+/// keeps a memo of its own keeps its own answers, so a quantity that depends
+/// on a kernel is remembered per kernel and per device.
+class DeviceMemo {
+ public:
+  template 
+  long get(const Query& query) {
+    int device = 0;
+    check_query(cudaGetDevice(&device), "the device query");
+    const std::lock_guard lock(mutex_);
+    const auto known = values_.find(device);
+    if (known != values_.end()) {
+      return known->second;
+    }
+    const long value = query(device);
+    values_.emplace(device, value);
+    return value;
+  }
+
+ private:
+  std::mutex mutex_;
+  std::map values_;
+};
+
+/// Multiprocessors of `device`.
+inline long multiprocessor_count(int device) {
+  int multiprocessors = 0;
+  check_query(cudaDeviceGetAttribute(&multiprocessors,
+                                     cudaDevAttrMultiProcessorCount, device),
+              "the multiprocessor count");
+  return multiprocessors;
+}
+
+/// Warps that share one destination row: the most among `1, 2, 4, 8` whose
+/// warps fit in one wave of the warps that `Kernel`, the scan of blocks of
+/// `Rows` independent rows that serves the system otherwise, keeps resident
+/// across the current device.
+template 
+int row_split(long node_count) {
+  static DeviceMemo capacities;
+  const long capacity = capacities.get([](int device) {
+    int resident = 0;
+    check_query(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+                    &resident, Kernel, Rows * kWarpSize, 0),
+                "the occupancy query of an edge scan");
+    return multiprocessor_count(device) * resident * Rows;
+  });
+  int split = 1;
+  while (split < 8 && node_count * split * 2 <= capacity) {
+    split *= 2;
+  }
+  return split;
+}
+
+/// Whether a system fills the current device with co-scheduled blocks of
+/// `Warps` rows: one such block fills a multiprocessor, so a run with fewer
+/// than two blocks per multiprocessor would leave part of the device idle, and
+/// narrow blocks then serve it better.
+template 
+bool fills_coscheduled(long node_count) {
+  if constexpr (Warps == 1) {
+    return false;
+  } else {
+    static DeviceMemo capacities;
+    const long capacity = capacities.get(
+        [](int device) { return 2L * multiprocessor_count(device) * Warps; });
+    return node_count >= capacity;
+  }
+}
+
+/// The type every node of the system carries, or -1 when the scan has to
+/// read the ordered pair of every edge. The answer is uniform over the grid.
+__device__ __forceinline__ int single_type(const Arguments& args,
+                                           int radial_modes) {
+  if (args.type_mismatch == nullptr || radial_modes > kFoldedModes ||
+      __ldg(args.type_mismatch) != 0) {
+    return -1;
+  }
+  return static_cast(__ldg(args.atype));
+}
+
+/// Value of one extra radial feature of a pair-uniform row.
+__device__ __forceinline__ float extra_feature(const TableRow& row,
+                                               int channels,
+                                               int feature,
+                                               int radial_modes,
+                                               float coordinate) {
+  if (feature < radial_modes) {
+    return evaluate_table(row, channels + feature, coordinate);
+  }
+  return feature == radial_modes ? 1.0f : 0.0f;
+}
+
+/// Gather the nine reduced harmonic sums of one extra feature onto every
+/// lane. After the merge of the edge groups the lane of feature index `f`
+/// modulo the group width holds them in its tile `f / Width`.
+template 
+__device__ __forceinline__ void gather_extra_feature(
+    const float (&extra)[Tiles][9], int feature, float (&sums)[9]) {
+#pragma unroll
+  for (int component = 0; component < 9; ++component) {
+    float value = 0.0f;
+#pragma unroll
+    for (int tile = 0; tile < Tiles; ++tile) {
+      if (tile == feature / Width) {
+        value = extra[tile][component];
+      }
+    }
+    sums[component] = __shfl_sync(kWarpMask, value, feature % Width);
+  }
+}
+
+/// Expand the low-rank sums of a single-type row into the scalar sums of the
+/// channels from ``C1`` on (see `low_rank_basis`). On entry the tiles from
+/// ``C1 / Width`` of ``scalar`` hold the sums of the lane's basis columns, the
+/// same on every edge group after their merge; on exit every tile from
+/// ``C1 / Width`` holds the sum of the lane's channel, accumulated over the
+/// basis functions in increasing order.
+template 
+__device__ __forceinline__ void expand_basis(
+    float (&scalar)[Tiles],
+    const float* __restrict__ basis_matrix,
+    int base_channel) {
+  constexpr int Direct = P::C1 / Width;
+  constexpr int BasisTiles = P::Basis / Width;
+  constexpr int Expanded = P::C0 - P::C1;
+  float basis[BasisTiles];
+#pragma unroll
+  for (int tile = 0; tile < BasisTiles; ++tile) {
+    basis[tile] = scalar[Direct + tile];
+  }
+  float sum[Tiles - Direct] = {};
+#pragma unroll
+  for (int k = 0; k < P::Basis; ++k) {
+    const float value = __shfl_sync(kWarpMask, basis[k / Width], k % Width);
+    const float* row = basis_matrix + k * Expanded + base_channel;
+#pragma unroll
+    for (int tile = 0; tile < Tiles - Direct; ++tile) {
+      sum[tile] = fmaf(__ldg(row + tile * Width), value, sum[tile]);
+    }
+  }
+#pragma unroll
+  for (int tile = 0; tile < Tiles - Direct; ++tile) {
+    scalar[Direct + tile] = sum[tile];
+  }
+}
+
+/// Stage the expansion matrix of the low-rank basis transposed into the
+/// shared `transposed` with the threads of the block: coefficient `k` of
+/// channel `C1 + c` at `c * (Basis + 1) + k`. The lanes of a warp that each
+/// project the cotangents of a row onto one basis function then read
+/// consecutive words, where the matrix itself places them a row of `C0 - C1`
+/// coefficients apart and every load touches a cache line per lane; the
+/// padded rows keep the transposing stores free of bank conflicts.
+template 
+__device__ __forceinline__ void stage_transposed_basis(
+    const float* __restrict__ basis_matrix, float* transposed) {
+  constexpr int Expanded = P::C0 - P::C1;
+  for (int index = threadIdx.x; index < P::Basis * Expanded;
+       index += blockDim.x) {
+    const int k = index / Expanded;
+    const int channel = index - k * Expanded;
+    transposed[channel * (P::Basis + 1) + k] = __ldg(basis_matrix + index);
+  }
+}
+
+/// Expansion coefficients `k` of the four channels from `C1 + first` on of
+/// the low-rank basis, for `first` a multiple of four, in one 16-byte load:
+/// a lane that projects onto one basis function reads its own row of the
+/// matrix, so a load touches a cache line per lane whatever its width, and
+/// the wide load serves four channels with it.
+template 
+__device__ __forceinline__ float4
+basis_quad(const float* __restrict__ basis_matrix, int k, int first) {
+  static_assert((P::C0 - P::C1) % 4 == 0,
+                "the rows of the expansion hold whole quads of channels");
+  return __ldg(reinterpret_cast(basis_matrix +
+                                               k * (P::C0 - P::C1) + first));
+}
+
+/// The coefficients of `basis_quad` from the transpose that the block staged
+/// in shared memory (`stage_transposed_basis`).
+template 
+__device__ __forceinline__ float4 staged_basis_quad(const float* transposed,
+                                                    int k,
+                                                    int first) {
+  constexpr int Stride = P::Basis + 1;
+  const float* column = transposed + first * Stride + k;
+  return make_float4(column[0], column[Stride], column[2 * Stride],
+                     column[3 * Stride]);
+}
+
+/// Project the scalar cotangents of the channels from ``C1`` on onto the
+/// low-rank basis of a single-type row: the cotangent of basis function ``k``
+/// is the sum over those channels, in increasing order, of the channel's
+/// cotangent times its expansion coefficient ``k``.
+template 
+__device__ __forceinline__ void project_basis(
+    const float* scalar_gradient,
+    const float* __restrict__ basis_matrix,
+    float* basis_gradient,
+    int lane) {
+  for (int k = lane; k < P::Basis; k += kWarpSize) {
+    float sum = 0.0f;
+    for (int first = 0; first < P::C0 - P::C1; first += 4) {
+      const float4 quad = basis_quad
(basis_matrix, k, first);
+      const float coefficients[4] = {quad.x, quad.y, quad.z, quad.w};
+#pragma unroll
+      for (int offset = 0; offset < 4; ++offset) {
+        sum = fmaf(coefficients[offset],
+                   scalar_gradient[P::C1 + first + offset], sum);
+      }
+    }
+    basis_gradient[k] = sum;
+  }
+}
+
+// === Lockstep windows ===
+
+/// Squared true length of the bond stored at `position`.
+__device__ __forceinline__ float bond_length_squared(const float* edge_vec,
+                                                     long position) {
+  const float x = edge_vec[position * 3 + 0];
+  const float y = edge_vec[position * 3 + 1];
+  const float z = edge_vec[position * 3 + 2];
+  return x * x + y * y + z * z;
+}
+
+/// End of the window that starts at `position`: the first position from there
+/// on whose bond is at least as long as the window bound, or `end`. The warp
+/// examines thirty-two positions at a time and answers uniformly.
+__device__ __forceinline__ long window_end(long position,
+                                           long end,
+                                           float bound_squared,
+                                           const float* edge_vec,
+                                           int lane) {
+  for (;;) {
+    const long candidate = position + lane;
+    const bool inside =
+        candidate < end &&
+        bond_length_squared(edge_vec, candidate) < bound_squared;
+    const unsigned run = __ballot_sync(kWarpMask, inside);
+    const int length = run == kWarpMask ? kWarpSize : __ffs(~run) - 1;
+    position += length;
+    if (length < kWarpSize) {
+      return position;
+    }
+  }
+}
+
+/// Visit the positions of the destination row `[begin, end)` that the edge
+/// group `group` of the warp `part` of the row owns, every position exactly
+/// once.
+///
+/// A row shared by `Split` warps deals its positions in runs of `Groups`, one
+/// run per warp in turn. A block of several warps, one row each, walks the
+/// rows in windows of `kLockstepWindow` Å of bond length (see `Lockstep`): a
+/// warp visits the positions of its row below the bound of the current
+/// window, and the block crosses a barrier before the next window. A window
+/// ends at the first position at or beyond its bound, so the windows partition
+/// every row whatever its order, and the positions left after the last window
+/// are visited after it. The windows stop at the one that holds `reach`, the
+/// longest bond the radial table covers, so a bond beyond it, which no valid
+/// graph holds, is visited after the last window as well. A warp without a row
+/// passes empty bounds and joins the barriers.
+///
+/// The block walks the windows from the lowest that holds the first bond of
+/// one of its rows to the highest that holds the last bond of one of them. A
+/// warp visits the same positions in every window, whichever rows share its
+/// block, when its row is sorted by bond length, as the canonical graphs of
+/// LAMMPS and of the test builders are, and for any row while the table spans
+/// at most two windows, as it does below a cutoff of 8 Å; at a longer cutoff
+/// the windows of an unsorted row can follow from the rows that share its
+/// block.
+template 
+__device__ __forceinline__ void walk_row(long begin,
+                                         long end,
+                                         int group,
+                                         int lane,
+                                         int part,
+                                         const float* edge_vec,
+                                         float reach,
+                                         const Visit& visit) {
+  constexpr float width = kLockstepWindow;
+  if constexpr (Warps == 1) {
+    for (long position = begin + part * Groups + group; position < end;
+         position += Split * Groups) {
+      visit(position);
+    }
+  } else {
+    __shared__ int range[2];
+    if (threadIdx.x == 0) {
+      range[0] = INT_MAX;
+      range[1] = -1;
+    }
+    __syncthreads();
+    constexpr float inverse_width = 1.0f / width;
+    if (lane == 0 && end > begin) {
+      // The bound keeps the conversion in range for any length, an infinite
+      // one included.
+      const float last_window = reach * inverse_width;
+      const auto window_of = [&](long position) {
+        return static_cast(fminf(
+            sqrtf(bond_length_squared(edge_vec, position)) * inverse_width,
+            last_window));
+      };
+      atomicMin(&range[0], window_of(begin));
+      atomicMax(&range[1], window_of(end - 1));
+    }
+    __syncthreads();
+    const int first = range[0];
+    const int last = range[1];
+    long start = begin;
+    for (int window = first; window <= last; ++window) {
+      const float bound = static_cast(window + 1) * width;
+      const long stop = window_end(start, end, bound * bound, edge_vec, lane);
+      for (long position = start + group; position < stop; position += Groups) {
+        visit(position);
+      }
+      start = stop;
+      __syncthreads();
+    }
+    for (long position = start + group; position < end; position += Groups) {
+      visit(position);
+    }
+  }
+}
+
+// === Split rows ===
+//
+// A system with fewer destination rows than the device holds resident warps
+// leaves most of it idle under one warp per row, and every warp then runs
+// its whole row as one latency-bound chain. The edge scans of such a system
+// share each row among `Split` warps of one block instead: every warp scans
+// its share of the edges into its own accumulators, the warps past the first
+// publish theirs to shared memory, and the first warp adds them in warp order
+// before it completes the row as a single warp would. The order of every sum
+// therefore depends on `Split` alone, which is a function of the node count
+// and of the occupancy the device reports for the kernel, and the result is
+// reproducible on a given device.
+
+/// Accumulators exchanged per round of a split-row merge. The exchange
+/// buffer holds one round, so that it stays small next to the L1 capacity the
+/// table-bound scans depend on.
+constexpr int kSplitChunk = 8;
+
+/// Shared floats of the exchange buffer of a row shared by `Split` warps.
+template 
+constexpr int split_buffer_floats() {
+  return Split > 1 ? (Split - 1) * kSplitChunk * kWarpSize : 1;
+}
+
+/// Add the accumulators of the other warps of a shared row onto the first.
+///
+/// Every warp passes the same visitor, which presents its accumulators in one
+/// fixed order. In rounds of `kSplitChunk` accumulators, the warps past the
+/// first store theirs and the first adds them in warp order. The function
+/// returns whether the calling warp carries on with the row.
+template 
+__device__ __forceinline__ bool merge_split_row(int part,
+                                                int lane,
+                                                float* buffer,
+                                                const Visit& accumulators) {
+  if constexpr (Split == 1) {
+    return true;
+  } else {
+#pragma unroll
+    for (int first = 0; first < Count; first += kSplitChunk) {
+      if (part > 0) {
+        int slot = 0;
+        accumulators([&](float& value) {
+          if (slot >= first && slot < first + kSplitChunk) {
+            buffer[((part - 1) * kSplitChunk + slot - first) * kWarpSize +
+                   lane] = value;
+          }
+          ++slot;
+        });
+      }
+      __syncthreads();
+      if (part == 0) {
+        for (int other = 1; other < Split; ++other) {
+          int slot = 0;
+          accumulators([&](float& value) {
+            if (slot >= first && slot < first + kSplitChunk) {
+              value += buffer[((other - 1) * kSplitChunk + slot - first) *
+                                  kWarpSize +
+                              lane];
+            }
+            ++slot;
+          });
+        }
+      }
+      __syncthreads();
+    }
+    return part == 0;
+  }
+}
+
+/// Synchronize the warps that share a row.
+template 
+__device__ __forceinline__ void sync_row() {
+  if constexpr (Split > 1) {
+    __syncthreads();
+  } else {
+    __syncwarp();
+  }
+}
+
 // === Forward ===
 
 template 
+          typename index_t,
+          int Warps,
+          int Split,
+          Scan Fold>
 __global__ __launch_bounds__(
-    Profile::Threads,
-    edge_resident_blocks()) void forward_kernel(Arguments args) {
+    Warps * Split * kWarpSize,
+    (slots_types(Fold)
+         ? slot_resident_warps() ? Warps : 0>()
+         : edge_resident_blocks()) /
+        (Warps * Split)) void forward_kernel(Arguments args) {
   using P = Profile;
+  enter_edge_scan(
+      args);
+  constexpr bool SlotScan = slots_types(Fold);
   constexpr int EdgeWidth = P::ForwardEdgeWidth;
   constexpr int Groups = kWarpSize / EdgeWidth;
   constexpr int ChannelTiles = Channels / EdgeWidth;
   constexpr int AngularTiles = (P::C1 + EdgeWidth - 1) / EdgeWidth;
   constexpr int TensorTiles = (P::C2 + EdgeWidth - 1) / EdgeWidth;
   constexpr int HighTiles = (P::HighCount + EdgeWidth - 1) / EdgeWidth;
+  static_assert(
+      P::Basis == 0 || (P::C1 % EdgeWidth == 0 && P::Basis % EdgeWidth == 0),
+      "the low-rank columns tile the lanes of an edge group");
   // The spin families share the neighbour width with degree two, so they
   // inherit its tiling. The single neighbour quadrupole distributes its five
   // components across lanes, exactly as the single-channel high degrees do.
   constexpr int SpinTiles = HasSpin ? TensorTiles : 0;
   constexpr int SpinTensorTiles = HasSpin ? (5 + EdgeWidth - 1) / EdgeWidth : 0;
 
-  const int thread = threadIdx.x;
-  const long node = blockIdx.x;
-  if (node >= args.node_count) {
-    return;
-  }
+  static_assert(Warps == 1 || Split == 1,
+                "a block either holds several rows or shares one");
+  const int thread = threadIdx.x % kWarpSize;
+  // A block holds `Warps` rows, one warp each, or one row shared by `Split`
+  // warps; `part` is the warp's share of a shared row.
+  const int warp = Warps > 1 ? threadIdx.x / kWarpSize : 0;
+  const int part = Split > 1 ? threadIdx.x / kWarpSize : 0;
+  // Blocks of independent rows of the profiles up to degree two take their
+  // rows from the last block to the first, and every other block takes its
+  // rows in ascending order. The LAMMPS graph builder writes the rows in
+  // ascending order, so when one run covers the system, the rows it writes
+  // last are still in L2 when the forward scan starts.
+  constexpr bool Windowed = windowed();
+  constexpr int WalkWarps = Windowed ? Warps : 1;
+  constexpr bool Reversed = !Windowed && Lmax <= 2;
+  const long block = Reversed ? static_cast(gridDim.x - 1 - blockIdx.x)
+                              : static_cast(blockIdx.x);
+  const long node = block * Warps + warp;
+  // A warp past the last node of the final block owns no row: it reads node
+  // zero for its constants and scans an empty row, so that it still reaches
+  // every barrier of its block.
+  const bool active = node < args.node_count;
+  const long row_node = active ? node : 0;
   const int base_channel = thread & (EdgeWidth - 1);
   const int group = thread / EdgeWidth;
   const int leader = group * EdgeWidth;
@@ -113,13 +1038,16 @@ __global__ __launch_bounds__(
   const auto* edge_index = static_cast(args.edge_index);
   const auto* destination_order =
       static_cast(args.destination_order);
-  const int center_type = static_cast(args.atype[args.node_begin + node]);
-  // The centre half of every pair length scale is a block constant, like the
-  // centre type it follows.
+  const int center_type =
+      static_cast(args.atype[args.node_begin + row_node]);
+  // The centre half of every pair length scale is a node constant, like the
+  // centre type it follows, and so is the length from which the window of no
+  // pair of the row needs resolving.
   const float center_radius =
       Bridged ? __ldg(args.contact_radius + center_type) : 0.0f;
-  const long begin = args.destination_row_ptr[node];
-  const long end = args.destination_row_ptr[node + 1];
+  const float bound = Bridged ? window_bound(args, center_radius) : 0.0f;
+  const long begin = args.destination_row_ptr[row_node];
+  const long end = active ? args.destination_row_ptr[row_node + 1] : begin;
   const int radial_modes = HasModes ? args.radial_modes : 0;
 
   float scalar[ChannelTiles] = {};
@@ -138,14 +1066,75 @@ __global__ __launch_bounds__(
   // a preceding shared array of arbitrary length would otherwise leave the
   // block on a four-byte boundary.
   __shared__ __align__(
-      16) float mode_cache[HasModes ? Groups * kModeStride : 1];
-  float* modes = mode_cache + (HasModes ? group * kModeStride : 0);
+      16) float mode_cache[Warps * Split][HasModes ? Groups * kModeStride : 1];
+  float* modes = mode_cache[warp + part] + (HasModes ? group * kModeStride : 0);
+
+  // A single-type system accumulates the bare radial features in the moment
+  // registers and the extra features in `extra`, laid out as one scalar,
+  // three vector and five tensor sums per feature tile.
+  constexpr bool Foldable = foldable();
+  constexpr int ExtraFeatures = (HasModes ? kFoldedModes : 0) + 1;
+  constexpr int ExtraTiles = (ExtraFeatures + EdgeWidth - 1) / EdgeWidth;
+  float extra[ExtraTiles][9] = {};
+  int fold_type = -1;
+  if constexpr (Foldable && !SlotScan) {
+    fold_type = single_type(args, radial_modes);
+  }
+  const bool folded = fold_type >= 0;
+  // A slotted scan accumulates the mode profiles and the basis functions of
+  // the edges of each node type of its system in the slot of that type.
+  static_assert(
+      !SlotScan || slottable(),
+      "a slotted scan serves a profile that slots");
+  constexpr bool SlotModes = SlotScan && HasModes;
+  constexpr bool SlotBasis = slots_basis(Fold);
+  constexpr int Slots = type_slots(Fold);
+  constexpr int SlotTiles = SlotModes ? slot_tiles() : 1;
+  [[maybe_unused]] constexpr int SlotLanes = EdgeWidth / kFoldedModes;
+  static_assert(
+      !SlotModes || (EdgeWidth % kFoldedModes == 0 && Slots % SlotLanes == 0),
+      "the modes of the slots tile the lanes of an edge group");
+  [[maybe_unused]] float slot_extra[SlotTiles][9] = {};
+  // The channel tiles that a scan accumulates edge by edge: every tile, or
+  // for a slotted low-rank scan those below `C1`, whose FiLM stays on the
+  // edges; the slots of the basis functions provide the others per row.
+  constexpr int ScanTiles = SlotBasis ? P::C1 / EdgeWidth : ChannelTiles;
+  constexpr int BasisTiles = SlotBasis ? P::Basis / EdgeWidth : 1;
+  static_assert(!SlotBasis || (Groups <= Slots && Slots % Groups == 0 &&
+                               EdgeWidth >= Slots),
+                "the edge groups share the slots of the basis expansion, and "
+                "the lanes of a group hold the envelope sums of the slots");
+  [[maybe_unused]] float basis_slot[BasisTiles][Slots] = {};
+  [[maybe_unused]] float slot_mass = 0.0f;
+  TypeSlots slots;
+  if constexpr (SlotScan) {
+    slots = load_type_slots(args);
+  }
+  // The high degrees read the full amplitude of the leading channel, whose
+  // coefficients a folded scan hoists.
+  const long fixed_pair =
+      static_cast(center_type) * args.type_count + fold_type;
+  const float2 leading_film =
+      folded ? __ldg(reinterpret_cast(args.pair_film) +
+                     fixed_pair * Channels)
+             : make_float2(0.0f, 0.0f);
+  const float* leading_mixing =
+      HasModes && folded
+          ? args.pair_mixing + fixed_pair * Channels * radial_modes
+          : nullptr;
 
   // === Step 1. Reduce the destination row into degree-wise moments ===
-  for (long position = begin + group; position < end; position += Groups) {
+  const auto scan_edge = [&](const long position, auto scan) {
+    constexpr Scan Kind = decltype(scan)::value;
+    constexpr bool Folded = Kind == Scan::Direct || Kind == Scan::LowRank;
+    constexpr int FoldColumns = fold_columns();
+    constexpr int FoldTiles = FoldColumns / EdgeWidth;
+    // Channel tiles whose FiLM amplitude the edges carry.
+    constexpr int EdgeTiles =
+        slots_basis(Kind) ? P::C1 / EdgeWidth : ChannelTiles;
     const long edge = edge_at_position(position, destination_order);
     if (args.edge_mask != nullptr && !args.edge_mask[edge]) {
-      continue;
+      return;
     }
     // Every lane of the group reloads the shared edge state instead of
     // broadcasting it from a leader. The addresses are identical inside the
@@ -156,168 +1145,423 @@ __global__ __launch_bounds__(
     if constexpr (!Canonical) {
       if (center_type >= args.type_count - 1 ||
           source_type >= args.type_count - 1) {
-        continue;
+        return;
       }
     }
     const EdgeGeometry geometry = load_geometry(
         edge, source, source_type, args.rcut, args.f_inner, args.f_outer,
-        args.inverse_f_width, center_radius, args.contact_radius, args.eps,
-        args.edge_vec);
+        args.inverse_f_width, center_radius, bound, args.contact_radius,
+        args.eps, args.edge_vec);
     const TableLocation location = locate_table(
         geometry.seen, args.table_stride, args.table_max, args.interval_count);
-    const TableRow row = table_row(args.table, location, args.table_width);
+    const TableRow row =
+        reads_low_rank(Kind)
+            ? table_row(args.basis_table, location, FoldColumns + radial_modes)
+            : table_row(args.table, location, args.table_width);
     const float coordinate = location.coordinate;
-    if constexpr (HasModes) {
-      __syncwarp(mask);
-      for (int mode = base_channel; mode < radial_modes; mode += EdgeWidth) {
-        modes[mode] = evaluate_table(row, Channels + mode, coordinate);
-      }
-      __syncwarp(mask);
-    }
-
     float basis[9];
     fill_angular_basis(geometry, basis);
     const float envelope = geometry.envelope;
-    const long pair =
-        static_cast(center_type) * args.type_count + geometry.source_type;
-    const float2* film_row =
-        reinterpret_cast(args.pair_film + pair * Channels * 2) +
-        base_channel;
-    const float* mixing_row =
-        HasModes
-            ? args.pair_mixing + (pair * Channels + base_channel) * radial_modes
-            : nullptr;
+    const float squared = envelope * envelope;
+    // Angular payload of the leading channel, held by the leader of the group,
+    // and for a slotted scan its amplitude before the mode mixing.
+    float leading_payload = 0.0f;
+    [[maybe_unused]] float leading_value = 0.0f;
 
-    // The spin payload is hoisted out of the channel loop: every channel
-    // multiplies the same conditioned neighbour moment and the same projection
-    // of that moment onto the bond, and only the amplitude varies with the
-    // channel.
-    float neighbor_spin[3] = {0.0f, 0.0f, 0.0f};
-    float neighbor_magnitude = 0.0f;
-    float neighbor_gate = 0.0f;
-    float bond_alignment = 0.0f;
-    const float2* spin_row = nullptr;
-    if constexpr (HasSpin) {
-      const SpinTypeWeights weights =
-          load_spin_type(args.spin_type, geometry.source_type);
-      load_conditioned_spin(args.spin, weights.scale, geometry.source,
-                            neighbor_spin);
-      neighbor_magnitude = neighbor_spin[0] * neighbor_spin[0] +
-                           neighbor_spin[1] * neighbor_spin[1] +
-                           neighbor_spin[2] * neighbor_spin[2];
-      neighbor_gate = weights.gate;
-      // Component of the neighbour moment along the bond. The bond-projected
-      // family is this scalar times the unit direction, which is `basis[1..3]`.
-      bond_alignment = neighbor_spin[0] * basis[1] +
-                       neighbor_spin[1] * basis[2] +
-                       neighbor_spin[2] * basis[3];
-      spin_row =
-          reinterpret_cast(args.spin_pair + pair * P::Cs * 2) +
-          base_channel;
-    }
-
-    float angular_zero = 0.0f;
-    float spin_zero = 0.0f;
+    if constexpr (Folded) {
+      float leading = 0.0f;
+      // A column below ``C1`` is the channel of that index; a low-rank table
+      // continues with the basis functions, which feed degree zero alone.
 #pragma unroll
-    for (int tile = 0; tile < ChannelTiles; ++tile) {
-      const int channel = base_channel + tile * EdgeWidth;
-      const float radial = evaluate_table(row, channel, coordinate);
-      const float2 film = __ldg(film_row + tile * EdgeWidth);
-      float film_value = fmaf(film.x, radial, film.y);
-      if constexpr (HasModes) {
-        accumulate_modes(mixing_row + tile * (EdgeWidth * radial_modes), modes,
-                         radial_modes, film_value);
-      }
-      // Degree zero carries one envelope factor; every non-scalar degree
-      // carries a second one.
-      const float amplitude = film_value * envelope;
-      const float angular = amplitude * envelope;
-      scalar[tile] += amplitude;
-      if (channel < P::C1) {
+      for (int tile = 0; tile < FoldTiles; ++tile) {
+        const int channel = base_channel + tile * EdgeWidth;
+        const float radial = evaluate_table(row, channel, coordinate);
+        scalar[tile] = fmaf(radial, envelope, scalar[tile]);
+        const float angular = radial * squared;
+        if (channel < P::C1) {
 #pragma unroll
-        for (int component = 0; component < 3; ++component) {
-          vector[tile][component] =
-              fmaf(angular, basis[1 + component], vector[tile][component]);
+          for (int component = 0; component < 3; ++component) {
+            vector[tile][component] =
+                fmaf(angular, basis[1 + component], vector[tile][component]);
+          }
+        }
+        if (channel < P::C2) {
+#pragma unroll
+          for (int component = 0; component < 5; ++component) {
+            tensor[tile][component] =
+                fmaf(angular, basis[4 + component], tensor[tile][component]);
+          }
+        }
+        if (tile == 0) {
+          leading = radial;
         }
       }
-      if (channel < P::C2) {
+      float feature_value[ExtraTiles] = {};
 #pragma unroll
-        for (int component = 0; component < 5; ++component) {
-          tensor[tile][component] =
-              fmaf(angular, basis[4 + component], tensor[tile][component]);
+      for (int tile = 0; tile < ExtraTiles; ++tile) {
+        if (tile * EdgeWidth <= radial_modes) {
+          const float value =
+              extra_feature(row, FoldColumns, base_channel + tile * EdgeWidth,
+                            radial_modes, coordinate);
+          feature_value[tile] = value;
+          extra[tile][0] = fmaf(value, envelope, extra[tile][0]);
+          const float angular = value * squared;
+#pragma unroll
+          for (int component = 1; component < 9; ++component) {
+            extra[tile][component] =
+                fmaf(angular, basis[component], extra[tile][component]);
+          }
         }
       }
+      if constexpr (Lmax >= 3) {
+        float film_value = fmaf(leading_film.x, leading, leading_film.y);
+        if constexpr (HasModes) {
+#pragma unroll
+          for (int mode = 0; mode < kFoldedModes; ++mode) {
+            if (mode < radial_modes) {
+              const float profile =
+                  __shfl_sync(mask, feature_value[mode / EdgeWidth],
+                              leader + mode % EdgeWidth);
+              film_value =
+                  fmaf(__ldg(leading_mixing + mode), profile, film_value);
+            }
+          }
+        }
+        leading_payload = film_value * squared;
+      }
+    } else {
+      if constexpr (HasModes && Kind == Scan::Pairwise) {
+        __syncwarp(mask);
+        for (int mode = base_channel; mode < radial_modes; mode += EdgeWidth) {
+          modes[mode] = evaluate_table(row, Channels + mode, coordinate);
+        }
+        __syncwarp(mask);
+      }
+      const long pair = static_cast(center_type) * args.type_count +
+                        geometry.source_type;
+      const float2* film_row = reinterpret_cast(
+                                   args.pair_film + pair * Channels * 2) +
+                               base_channel;
+      [[maybe_unused]] const float* mixing_row =
+          HasModes ? args.pair_mixing +
+                         (pair * Channels + base_channel) * radial_modes
+                   : nullptr;
+
+      // The spin payload is hoisted out of the channel loop: every channel
+      // multiplies the same conditioned neighbour moment and the same
+      // projection of that moment onto the bond, and only the amplitude varies
+      // with the channel.
+      [[maybe_unused]] float neighbor_spin[3] = {0.0f, 0.0f, 0.0f};
+      [[maybe_unused]] float neighbor_magnitude = 0.0f;
+      [[maybe_unused]] float neighbor_gate = 0.0f;
+      [[maybe_unused]] float bond_alignment = 0.0f;
+      [[maybe_unused]] const float2* spin_row = nullptr;
       if constexpr (HasSpin) {
-        if (channel < P::Cs) {
-          const float2 spin_film = __ldg(spin_row + tile * EdgeWidth);
-          const float weight =
-              fmaf(spin_film.x, radial, spin_film.y) * envelope * envelope;
-          spin_magnitude[tile] =
-              fmaf(weight, neighbor_magnitude, spin_magnitude[tile]);
-          spin_coordination[tile] =
-              fmaf(weight, neighbor_gate, spin_coordination[tile]);
-          const float projected = weight * bond_alignment;
+        const SpinTypeWeights weights =
+            load_spin_type(args.spin_type, geometry.source_type);
+        load_conditioned_spin(args.spin, weights.scale, geometry.source,
+                              neighbor_spin);
+        neighbor_magnitude = neighbor_spin[0] * neighbor_spin[0] +
+                             neighbor_spin[1] * neighbor_spin[1] +
+                             neighbor_spin[2] * neighbor_spin[2];
+        neighbor_gate = weights.gate;
+        // Component of the neighbour moment along the bond. The bond-projected
+        // family is this scalar times the unit direction, which is
+        // `basis[1..3]`.
+        bond_alignment = neighbor_spin[0] * basis[1] +
+                         neighbor_spin[1] * basis[2] +
+                         neighbor_spin[2] * basis[3];
+        spin_row =
+            reinterpret_cast(args.spin_pair + pair * P::Cs * 2) +
+            base_channel;
+      }
+
+      [[maybe_unused]] float spin_zero = 0.0f;
+#pragma unroll
+      for (int tile = 0; tile < EdgeTiles; ++tile) {
+        const int channel = base_channel + tile * EdgeWidth;
+        const float radial = evaluate_table(row, channel, coordinate);
+        const float2 film = __ldg(film_row + tile * EdgeWidth);
+        float film_value = fmaf(film.x, radial, film.y);
+        if constexpr (HasModes && Kind == Scan::Pairwise) {
+          accumulate_modes(mixing_row + tile * (EdgeWidth * radial_modes),
+                           modes, radial_modes, film_value);
+        }
+        // Degree zero carries one envelope factor; every non-scalar degree
+        // carries a second one.
+        const float amplitude = film_value * envelope;
+        const float angular = amplitude * envelope;
+        scalar[tile] += amplitude;
+        if (channel < P::C1) {
 #pragma unroll
           for (int component = 0; component < 3; ++component) {
-            spin_vector[tile][component] = fmaf(
-                weight, neighbor_spin[component], spin_vector[tile][component]);
-            spin_bond[tile][component] = fmaf(projected, basis[1 + component],
-                                              spin_bond[tile][component]);
+            vector[tile][component] =
+                fmaf(angular, basis[1 + component], vector[tile][component]);
           }
-          if (channel == 0) {
-            spin_zero = weight;
+        }
+        if (channel < P::C2) {
+#pragma unroll
+          for (int component = 0; component < 5; ++component) {
+            tensor[tile][component] =
+                fmaf(angular, basis[4 + component], tensor[tile][component]);
+          }
+        }
+        if constexpr (HasSpin) {
+          if (channel < P::Cs) {
+            const float2 spin_film = __ldg(spin_row + tile * EdgeWidth);
+            const float weight =
+                fmaf(spin_film.x, radial, spin_film.y) * envelope * envelope;
+            spin_magnitude[tile] =
+                fmaf(weight, neighbor_magnitude, spin_magnitude[tile]);
+            spin_coordination[tile] =
+                fmaf(weight, neighbor_gate, spin_coordination[tile]);
+            const float projected = weight * bond_alignment;
+#pragma unroll
+            for (int component = 0; component < 3; ++component) {
+              spin_vector[tile][component] =
+                  fmaf(weight, neighbor_spin[component],
+                       spin_vector[tile][component]);
+              spin_bond[tile][component] = fmaf(projected, basis[1 + component],
+                                                spin_bond[tile][component]);
+            }
+            if (channel == 0) {
+              spin_zero = weight;
+            }
+          }
+        }
+        if (channel == 0) {
+          leading_payload = angular;
+          leading_value = film_value;
+        }
+      }
+      if constexpr (HasSpin) {
+        // The single neighbour quadrupole reads the leading spin channel only,
+        // so its five components are distributed across the lanes of the edge
+        // group and each is evaluated on the one lane that owns it.
+        const float weight = __shfl_sync(mask, spin_zero, leader);
+#pragma unroll
+        for (int component = 0; component < 5; ++component) {
+          if (component % EdgeWidth == base_channel) {
+            spin_tensor[component / EdgeWidth] = fmaf(
+                weight,
+                spin_quadrupole_component(neighbor_spin[0], neighbor_spin[1],
+                                          neighbor_spin[2], component),
+                spin_tensor[component / EdgeWidth]);
           }
         }
       }
-      if constexpr (Lmax >= 3) {
-        if (channel == 0) {
-          angular_zero = angular;
-        }
+      if constexpr (slots_types(Kind)) {
+        const int slot = type_slot(slots, geometry.source_type);
+        if constexpr (HasModes) {
+          // The bare profile of the mode of the lane enters the slot of the
+          // source type, on the lanes that own that slot.
+          const int mode = base_channel % kFoldedModes;
+          const float profile =
+              mode < radial_modes
+                  ? evaluate_table(row, FoldColumns + mode, coordinate)
+                  : 0.0f;
+          const float angular = profile * squared;
+#pragma unroll
+          for (int tile = 0; tile < SlotTiles; ++tile) {
+            if (owned_slot(base_channel, tile) == slot) {
+              slot_extra[tile][0] =
+                  fmaf(profile, envelope, slot_extra[tile][0]);
+#pragma unroll
+              for (int component = 1; component < 9; ++component) {
+                slot_extra[tile][component] = fmaf(angular, basis[component],
+                                                   slot_extra[tile][component]);
+              }
+            }
+          }
+          if constexpr (Lmax >= 3) {
+            // The high degrees read the full amplitude of the leading
+            // channel, which the leader completes with the mode mixing of
+            // the pair; lane `leader + r` of the group holds the profile of
+            // mode `r`.
+            const float* pair_modes =
+                args.pair_mixing + pair * Channels * radial_modes;
+            float value = leading_value;
+#pragma unroll
+            for (int mixed = 0; mixed < kFoldedModes; ++mixed) {
+              if (mixed < radial_modes) {
+                value = fmaf(__ldg(pair_modes + mixed),
+                             __shfl_sync(mask, profile, leader + mixed), value);
+              }
+            }
+            leading_payload = value * squared;
+          }
+        }
+        if constexpr (slots_basis(Kind)) {
+          // The bare basis functions of the lane enter the slot of the source
+          // type, and the envelope enters it on the lane of that slot. Every
+          // slot takes the product with its weight, the envelope for the
+          // slot of the edge and zero for the others, which adds the exact
+          // zero and keeps the accumulation free of branches.
+          float weight[Slots];
+#pragma unroll
+          for (int entry = 0; entry < Slots; ++entry) {
+            weight[entry] = entry == slot ? envelope : 0.0f;
+          }
+#pragma unroll
+          for (int tile = 0; tile < BasisTiles; ++tile) {
+            const float radial = evaluate_table(
+                row, P::C1 + base_channel + tile * EdgeWidth, coordinate);
+#pragma unroll
+            for (int entry = 0; entry < Slots; ++entry) {
+              basis_slot[tile][entry] =
+                  fmaf(radial, weight[entry], basis_slot[tile][entry]);
+            }
+          }
+          slot_mass += base_channel == slot ? envelope : 0.0f;
+        }
+      }
+    }
+    if constexpr (Lmax >= 3) {
+      // Degrees three and above read only the leading channel, so their
+      // components are distributed across the lanes of the edge group. The
+      // component is iterated at compile time and selected with a lane
+      // predicate, which folds the harmonic selector and keeps the tile index
+      // constant.
+      const float amplitude = __shfl_sync(mask, leading_payload, leader);
+#pragma unroll
+      for (int component = 0; component < P::HighCount; ++component) {
+        if (component % EdgeWidth == base_channel) {
+          high[component / EdgeWidth] =
+              fmaf(amplitude, high_basis_value(geometry, component),
+                   high[component / EdgeWidth]);
+        }
+      }
+    }
+    if (base_channel == 0) {
+      scalar_mass += squared;
+      angular_mass = fmaf(squared, squared, angular_mass);
+    }
+  };
+  const auto pairwise_edge = [&](const long position) {
+    scan_edge(position, ScanTag{});
+  };
+  // The first warp of a shared row requests the lines of the row for every
+  // warp of the row.
+  if constexpr (Canonical) {
+    if (part == 0) {
+      prefetch_row_bonds(args, begin, end, thread);
+      if constexpr (SlotScan) {
+        prefetch_row_sources(args, edge_index, begin, end, thread);
+      }
+    }
+  }
+  if constexpr (SlotScan) {
+    walk_row(
+        begin, end, group, thread, part, args.edge_vec, args.table_max,
+        [&](const long position) { scan_edge(position, ScanTag{}); });
+  } else if constexpr (Foldable) {
+    if (folded) {
+      walk_row(
+          begin, end, group, thread, part, args.edge_vec, args.table_max,
+          [&](const long position) { scan_edge(position, ScanTag{}); });
+    } else {
+      walk_row(begin, end, group, thread, part,
+                                         args.edge_vec, args.table_max,
+                                         pairwise_edge);
+    }
+  } else {
+    walk_row(begin, end, group, thread, part,
+                                       args.edge_vec, args.table_max,
+                                       pairwise_edge);
+  }
+  if (!active) {
+    return;
+  }
+  // === Step 1b. Merge the warps that share the row ===
+  // The accumulators that `accumulators` visits, in its order.
+  constexpr int Accumulators = ScanTiles + 3 * AngularTiles + 5 * TensorTiles +
+                               (Lmax >= 3 ? HighTiles : 0) + 2 + 8 * SpinTiles +
+                               SpinTensorTiles +
+                               (Foldable && !SlotScan ? 9 * ExtraTiles : 0) +
+                               (SlotModes ? 9 * SlotTiles : 0) +
+                               (SlotBasis ? BasisTiles * Slots + 1 : 0);
+  __shared__ float split_buffer[split_buffer_floats()];
+  const auto accumulators = [&](auto&& visit) {
+#pragma unroll
+    for (int tile = 0; tile < ScanTiles; ++tile) {
+      visit(scalar[tile]);
+    }
+#pragma unroll
+    for (int tile = 0; tile < AngularTiles; ++tile) {
+#pragma unroll
+      for (int component = 0; component < 3; ++component) {
+        visit(vector[tile][component]);
+      }
+    }
+#pragma unroll
+    for (int tile = 0; tile < TensorTiles; ++tile) {
+#pragma unroll
+      for (int component = 0; component < 5; ++component) {
+        visit(tensor[tile][component]);
+      }
+    }
+    if constexpr (Lmax >= 3) {
+#pragma unroll
+      for (int tile = 0; tile < HighTiles; ++tile) {
+        visit(high[tile]);
+      }
+    }
+    visit(scalar_mass);
+    visit(angular_mass);
+    if constexpr (HasSpin) {
+#pragma unroll
+      for (int tile = 0; tile < SpinTiles; ++tile) {
+        visit(spin_magnitude[tile]);
+        visit(spin_coordination[tile]);
+#pragma unroll
+        for (int component = 0; component < 3; ++component) {
+          visit(spin_vector[tile][component]);
+          visit(spin_bond[tile][component]);
+        }
+      }
+#pragma unroll
+      for (int tile = 0; tile < SpinTensorTiles; ++tile) {
+        visit(spin_tensor[tile]);
       }
     }
-    if constexpr (HasSpin) {
-      // The single neighbour quadrupole reads the leading spin channel only,
-      // so its five components are distributed across the lanes of the edge
-      // group and each is evaluated on the one lane that owns it.
-      const float weight = __shfl_sync(mask, spin_zero, leader);
+    if constexpr (Foldable && !SlotScan) {
 #pragma unroll
-      for (int component = 0; component < 5; ++component) {
-        if (component % EdgeWidth == base_channel) {
-          spin_tensor[component / EdgeWidth] =
-              fmaf(weight,
-                   spin_quadrupole_component(neighbor_spin[0], neighbor_spin[1],
-                                             neighbor_spin[2], component),
-                   spin_tensor[component / EdgeWidth]);
+      for (int tile = 0; tile < ExtraTiles; ++tile) {
+#pragma unroll
+        for (int component = 0; component < 9; ++component) {
+          visit(extra[tile][component]);
         }
       }
     }
-    if constexpr (Lmax >= 3) {
-      // Degrees three and above read only the leading channel, so their
-      // components are distributed across the lanes of the edge group. The
-      // component is iterated at compile time and selected with a lane
-      // predicate, which folds the harmonic selector and keeps the tile index
-      // constant.
-      const float amplitude = __shfl_sync(mask, angular_zero, leader);
+    if constexpr (SlotModes) {
 #pragma unroll
-      for (int component = 0; component < P::HighCount; ++component) {
-        if (component % EdgeWidth == base_channel) {
-          high[component / EdgeWidth] =
-              fmaf(amplitude, high_basis_value(geometry, component),
-                   high[component / EdgeWidth]);
+      for (int tile = 0; tile < SlotTiles; ++tile) {
+#pragma unroll
+        for (int component = 0; component < 9; ++component) {
+          visit(slot_extra[tile][component]);
         }
       }
     }
-    if (base_channel == 0) {
-      const float squared = envelope * envelope;
-      scalar_mass += squared;
-      angular_mass = fmaf(squared, squared, angular_mass);
+    if constexpr (SlotBasis) {
+#pragma unroll
+      for (int tile = 0; tile < BasisTiles; ++tile) {
+#pragma unroll
+        for (int entry = 0; entry < Slots; ++entry) {
+          visit(basis_slot[tile][entry]);
+        }
+      }
+      visit(slot_mass);
     }
+  };
+  if (!merge_split_row(part, thread, split_buffer,
+                                            accumulators)) {
+    return;
   }
 
   // === Step 2. Merge the concurrent edge groups and normalize ===
   if constexpr (Groups > 1) {
 #pragma unroll
-    for (int tile = 0; tile < ChannelTiles; ++tile) {
+    for (int tile = 0; tile < ScanTiles; ++tile) {
       scalar[tile] = reduce_channel_groups(scalar[tile]);
     }
 #pragma unroll
@@ -363,28 +1607,213 @@ __global__ __launch_bounds__(
       }
     }
   }
-  __shared__ float normalizer_shared[4];
-  {
-    const float total_scalar = warp_sum(scalar_mass);
-    const float total_angular = warp_sum(angular_mass);
-    if (thread == 0) {
-      normalizer_shared[0] = rsqrtf(total_scalar + args.degree_floor);
-      normalizer_shared[1] = rsqrtf(total_angular + args.degree_floor);
-      normalizer_shared[2] = sqrtf(total_scalar + args.degree_floor);
-      normalizer_shared[3] = sqrtf(total_angular + args.degree_floor);
+  if (folded) {
+    // After the merge of the edge groups, every lane holds the row totals of
+    // its channels, and the lane of each extra feature its sums. The
+    // coefficients of the pair turn the bare sums into the moments of the
+    // amplitude.
+#pragma unroll
+    for (int tile = 0; tile < ExtraTiles; ++tile) {
+#pragma unroll
+      for (int component = 0; component < 9; ++component) {
+        extra[tile][component] =
+            reduce_channel_groups(extra[tile][component]);
+      }
+    }
+    const long pair =
+        static_cast(center_type) * args.type_count + fold_type;
+    const float2* film_row =
+        reinterpret_cast(args.pair_film + pair * Channels * 2) +
+        base_channel;
+    if constexpr (Fold == Scan::LowRank) {
+      expand_basis(scalar, args.basis_matrix, base_channel);
+    }
+    float sums[9];
+    gather_extra_feature(extra, radial_modes, sums);
+#pragma unroll
+    for (int tile = 0; tile < ChannelTiles; ++tile) {
+      const int channel = base_channel + tile * EdgeWidth;
+      const float2 film = __ldg(film_row + tile * EdgeWidth);
+      scalar[tile] = fmaf(film.x, scalar[tile], film.y * sums[0]);
+      if (channel < P::C1) {
+#pragma unroll
+        for (int component = 0; component < 3; ++component) {
+          vector[tile][component] = fmaf(film.x, vector[tile][component],
+                                         film.y * sums[1 + component]);
+        }
+      }
+      if (channel < P::C2) {
+#pragma unroll
+        for (int component = 0; component < 5; ++component) {
+          tensor[tile][component] = fmaf(film.x, tensor[tile][component],
+                                         film.y * sums[4 + component]);
+        }
+      }
+    }
+    if constexpr (HasModes) {
+      const float* mixing_row =
+          args.pair_mixing + (pair * Channels + base_channel) * radial_modes;
+      for (int mode = 0; mode < radial_modes; ++mode) {
+        gather_extra_feature(extra, mode, sums);
+#pragma unroll
+        for (int tile = 0; tile < ChannelTiles; ++tile) {
+          const int channel = base_channel + tile * EdgeWidth;
+          const float weight =
+              __ldg(mixing_row + tile * EdgeWidth * radial_modes + mode);
+          scalar[tile] = fmaf(weight, sums[0], scalar[tile]);
+          if (channel < P::C1) {
+#pragma unroll
+            for (int component = 0; component < 3; ++component) {
+              vector[tile][component] =
+                  fmaf(weight, sums[1 + component], vector[tile][component]);
+            }
+          }
+          if (channel < P::C2) {
+#pragma unroll
+            for (int component = 0; component < 5; ++component) {
+              tensor[tile][component] =
+                  fmaf(weight, sums[4 + component], tensor[tile][component]);
+            }
+          }
+        }
+      }
     }
   }
-  __syncthreads();
-  const float scalar_norm = normalizer_shared[0];
-  const float angular_norm = normalizer_shared[1];
+  if constexpr (SlotBasis) {
+    // After the merge of the edge groups every group holds the basis sums of
+    // the slots. Edge group `g` expands those of slots `g + j * Groups` into
+    // the bare channel sums from `C1` on, in increasing order of the basis
+    // functions, as a single-type scan does, and the FiLM scale and shift of
+    // the pair of the slot turn them into its share of the moments; the merge
+    // of the groups then adds the shares of the slots.
+    constexpr int Direct = P::C1 / EdgeWidth;
+    constexpr int Expanded = P::C0 - P::C1;
+#pragma unroll
+    for (int tile = 0; tile < BasisTiles; ++tile) {
+#pragma unroll
+      for (int entry = 0; entry < Slots; ++entry) {
+        basis_slot[tile][entry] =
+            reduce_channel_groups(basis_slot[tile][entry]);
+      }
+    }
+    slot_mass = reduce_channel_groups(slot_mass);
+    float share[ChannelTiles - Direct] = {};
+#pragma unroll
+    for (int round = 0; round < Slots / Groups; ++round) {
+      const int slot = group + round * Groups;
+      float sums[BasisTiles];
+#pragma unroll
+      for (int tile = 0; tile < BasisTiles; ++tile) {
+        sums[tile] = basis_slot[tile][0];
+#pragma unroll
+        for (int entry = 1; entry < Slots; ++entry) {
+          if (entry == slot) {
+            sums[tile] = basis_slot[tile][entry];
+          }
+        }
+      }
+      const float mass = __shfl_sync(kWarpMask, slot_mass, leader + slot);
+      float expanded[ChannelTiles - Direct] = {};
+#pragma unroll
+      for (int k = 0; k < P::Basis; ++k) {
+        const float value =
+            __shfl_sync(kWarpMask, sums[k / EdgeWidth], leader + k % EdgeWidth);
+        const float* row = args.basis_matrix + k * Expanded + base_channel;
+#pragma unroll
+        for (int tile = 0; tile < ChannelTiles - Direct; ++tile) {
+          expanded[tile] =
+              fmaf(__ldg(row + tile * EdgeWidth), value, expanded[tile]);
+        }
+      }
+      if (slot < slots.count) {
+        const long pair = static_cast(center_type) * args.type_count +
+                          slot_type(slots, slot);
+        const float2* film_row = reinterpret_cast(
+                                     args.pair_film + pair * Channels * 2) +
+                                 base_channel;
+#pragma unroll
+        for (int tile = 0; tile < ChannelTiles - Direct; ++tile) {
+          const float2 film = __ldg(film_row + (Direct + tile) * EdgeWidth);
+          share[tile] += fmaf(film.x, expanded[tile], film.y * mass);
+        }
+      }
+    }
+#pragma unroll
+    for (int tile = 0; tile < ChannelTiles - Direct; ++tile) {
+      scalar[Direct + tile] = reduce_channel_groups(share[tile]);
+    }
+  }
+  if constexpr (SlotModes) {
+    // After the merge of the edge groups every group holds the moments of the
+    // slots, and the mode mixing of the pair of each slot turns them into the
+    // mode share of the moments of every channel.
+#pragma unroll
+    for (int tile = 0; tile < SlotTiles; ++tile) {
+#pragma unroll
+      for (int component = 0; component < 9; ++component) {
+        slot_extra[tile][component] =
+            reduce_channel_groups(slot_extra[tile][component]);
+      }
+    }
+#pragma unroll
+    for (int slot = 0; slot < Slots; ++slot) {
+      if (slot < slots.count) {
+        const long pair = static_cast(center_type) * args.type_count +
+                          slot_type(slots, slot);
+        const float* mixing_row =
+            args.pair_mixing + (pair * Channels + base_channel) * radial_modes;
+#pragma unroll
+        for (int mode = 0; mode < kFoldedModes; ++mode) {
+          if (mode < radial_modes) {
+            float sums[9];
+#pragma unroll
+            for (int component = 0; component < 9; ++component) {
+              sums[component] = __shfl_sync(
+                  kWarpMask, slot_extra[slot / SlotLanes][component],
+                  slot_owner(slot, mode));
+            }
+#pragma unroll
+            for (int tile = 0; tile < ChannelTiles; ++tile) {
+              const int channel = base_channel + tile * EdgeWidth;
+              const float weight =
+                  __ldg(mixing_row + tile * EdgeWidth * radial_modes + mode);
+              scalar[tile] = fmaf(weight, sums[0], scalar[tile]);
+              if (channel < P::C1) {
+#pragma unroll
+                for (int component = 0; component < 3; ++component) {
+                  vector[tile][component] = fmaf(weight, sums[1 + component],
+                                                 vector[tile][component]);
+                }
+              }
+              if (channel < P::C2) {
+#pragma unroll
+                for (int component = 0; component < 5; ++component) {
+                  tensor[tile][component] = fmaf(weight, sums[4 + component],
+                                                 tensor[tile][component]);
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+  // Both masses are broadcast from lane zero, so every lane forms the same
+  // normalizers.
+  const float scalar_total = warp_sum(scalar_mass) + args.degree_floor;
+  const float angular_total = warp_sum(angular_mass) + args.degree_floor;
+  const float scalar_norm = rsqrtf(scalar_total);
+  const float angular_norm = rsqrtf(angular_total);
   if (thread < 2) {
     store_descriptor(args.descriptor, args.output_mean, args.output_inv_std,
                      node, P::OutputWidth, P::OutputDivisor + thread,
-                     normalizer_shared[2 + thread]);
+                     sqrtf(thread == 0 ? scalar_total : angular_total));
   }
 
-  // === Step 3. Publish the normalized moments and the saved state ===
-  __shared__ float moments[P::MomentWidth];
+  // === Step 3. Publish the normalized moments as the saved state ===
+  // The lanes of every edge group hold the merged moments after the
+  // reduction, so the leading group stores them in the flat state layout.
+  float* const moments = args.state_out + node * P::StateWidth;
   if (thread < EdgeWidth) {
 #pragma unroll
     for (int tile = 0; tile < ChannelTiles; ++tile) {
@@ -475,19 +1904,37 @@ __global__ __launch_bounds__(
       }
     }
   }
-  __syncthreads();
-
-  const long state_offset = node * P::StateWidth;
-  for (int coordinate = thread; coordinate < P::MomentWidth;
-       coordinate += P::Threads) {
-    args.state_out[state_offset + coordinate] = moments[coordinate];
-  }
   if (thread == 0) {
-    args.state_out[state_offset + P::MomentWidth] = scalar_norm;
-    args.state_out[state_offset + P::MomentWidth + 1] = angular_norm;
+    moments[P::MomentWidth] = scalar_norm;
+    moments[P::MomentWidth + 1] = angular_norm;
+  }
+}
+
+// === Forward readout ===
+
+// The invariant readout of the nodes whose moments the edge scan published:
+// the degree alignment, the probe projections and every invariant block of
+// the descriptor except the two divisors, which the scan emits itself. One
+// warp owns one node, as in the scan. The readout is a launch of its own
+// because the radial table rows that bound the scan are served by L1, whose
+// capacity follows the carveout the scan's shared memory selects; the readout
+// buffers of a resident node would double that carveout at full occupancy.
+template 
+__global__ __launch_bounds__(kWarpSize) void forward_readout_kernel(
+    Arguments args) {
+  using P = Profile;
+  const int thread = threadIdx.x;
+  const long node = blockIdx.x;
+  if (node >= args.node_count) {
+    return;
   }
+  const int center_type = static_cast(args.atype[args.node_begin + node]);
+  __shared__ float moments[P::MomentWidth];
+  stage_shared(
+      moments, args.state_out + node * P::StateWidth, thread);
+  __syncthreads();
 
-  // === Step 4. Align and project the wide degrees ===
+  // === Step 1. Align and project the wide degrees ===
   __shared__ float aligned[P::AlignedWidth];
   if (thread < P::C1) {
 #pragma unroll
@@ -546,7 +1993,7 @@ __global__ __launch_bounds__(
   }
   __syncthreads();
 
-  // === Step 5. Emit the invariant blocks ===
+  // === Step 2. Emit the invariant blocks ===
   for (int channel = thread; channel < Channels; channel += P::Threads) {
     store_descriptor(args.descriptor, args.output_mean, args.output_inv_std,
                      node, P::OutputWidth, P::OutputScalar + channel,
@@ -648,6 +2095,8 @@ __global__ __launch_bounds__(
   }
 
   if constexpr (Lmax >= 3) {
+    // The terms of one output are dealt to the lanes of the warp, and one
+    // reduction per output completes it.
     for (int record = 0; record < args.coupling_count; ++record) {
       const int* meta = args.coupling_meta + record * 8;
       const int degree_1 = meta[0];
@@ -658,13 +2107,13 @@ __global__ __launch_bounds__(
       const int probe_begin = meta[5];
       const int probe_count = meta[6];
       const int coordinate = meta[7];
-      for (int output = thread; output < probe_count; output += P::Threads) {
+      for (int output = 0; output < probe_count; ++output) {
         const int selection = __ldg(args.coupling_entry + probe_begin + output);
         const int index_1 = selection & 0xFF;
         const int index_2 = (selection >> 8) & 0xFF;
         const int index_3 = (selection >> 16) & 0xFF;
         float value = 0.0f;
-        for (int term = 0; term < nonzero_count; ++term) {
+        for (int term = thread; term < nonzero_count; term += P::Threads) {
           const int components =
               __ldg(args.coupling_entry + nonzero_begin + term);
           const float weight =
@@ -677,10 +2126,13 @@ __global__ __launch_bounds__(
               probes, moments, degree_3, (components >> 16) & 0xFF, index_3);
           value = fmaf(weight * first * second, third, value);
         }
-        store_descriptor(
-            args.descriptor, args.output_mean, args.output_inv_std, node,
-            P::OutputWidth, coordinate + output,
-            value * __ldg(args.coupling_value + probe_begin + output));
+        value = warp_sum(value);
+        if (thread == 0) {
+          store_descriptor(
+              args.descriptor, args.output_mean, args.output_inv_std, node,
+              P::OutputWidth, coordinate + output,
+              value * __ldg(args.coupling_value + probe_begin + output));
+        }
       }
     }
   }
@@ -706,7 +2158,7 @@ __global__ __launch_bounds__(
                          product[2] * product[2]);
   }
 
-  // === Step 6. Emit the spin invariants of even spin order ===
+  // === Step 3. Emit the spin invariants of even spin order ===
   if constexpr (HasSpin) {
     for (int pair = thread; pair < P::SpinGramVector; pair += P::Threads) {
       int row, column;
@@ -845,70 +2297,102 @@ __device__ __forceinline__ void add_bis222_probe_gradient(
   }
 }
 
-// Scatter the sparse-coupling VJP of one probe slot into a shared scratch row.
-// The scratch is indexed by harmonic component, which a register array cannot
-// address without spilling, and it is private to the lane, so the reduction
-// stays deterministic.
+// Index of one coupling operand among the targets of the coupling VJP: the
+// probe layout for degrees one and two, followed by the single-channel high
+// degrees in the order of the moment layout.
 template 
-__device__ __forceinline__ void accumulate_coupling_gradient(
+__device__ __forceinline__ int coupling_target(int degree,
+                                               int component,
+                                               int rank_index) {
+  using P = Profile;
+  if (degree == 1) {
+    return component * P::K1 + rank_index;
+  }
+  if (degree == 2) {
+    return 3 * P::K1 + component * P::K2 + rank_index;
+  }
+  return P::ProbeWidth + (degree == 3 ? 0 : P::High3) + component;
+}
+
+// VJP of every sparse coupling of one node, evaluated by the lanes of its node
+// group together. Each (output, term) pair of a record contributes to the
+// three operands it reads, and the pairs are dealt to the lanes in a fixed
+// order. Every lane accumulates into a private shared-memory row indexed by
+// `coupling_target`, and the rows are folded in lane order into `total`, so
+// the result is deterministic while no lane walks more than its share of the
+// pairs. Equal degrees occupy several operand slots of one record and receive
+// one contribution per slot.
+template 
+__device__ __forceinline__ void coupling_gradient(
+    int lane,
     long node,
     const Arguments& args,
     const float* __restrict__ probes,
     const float* __restrict__ moments,
-    int degree,
-    int rank_index,
-    float* __restrict__ scratch) {
-  using P = Profile;
+    float* __restrict__ rows,
+    float* __restrict__ total) {
+  using P = Profile;
+  constexpr int Width = P::NodeWidth;
+  constexpr int Targets = P::ProbeWidth + P::HighCount;
+  float* row = rows + lane * Targets;
+#pragma unroll
+  for (int target = 0; target < Targets; ++target) {
+    row[target] = 0.0f;
+  }
   for (int record = 0; record < args.coupling_count; ++record) {
     const int* meta = args.coupling_meta + record * 8;
     const int degrees[3] = {meta[0], meta[1], meta[2]};
-    if (degrees[0] != degree && degrees[1] != degree && degrees[2] != degree) {
-      continue;
-    }
     const int nonzero_begin = meta[3];
     const int nonzero_count = meta[4];
     const int probe_begin = meta[5];
     const int probe_count = meta[6];
     const int coordinate = meta[7];
-    for (int output = 0; output < probe_count; ++output) {
+    const int pairs = probe_count * nonzero_count;
+    if (pairs == 0) {
+      continue;
+    }
+    int output = lane / nonzero_count;
+    int term = lane - output * nonzero_count;
+    for (int pair = lane; pair < pairs; pair += Width) {
       const int selection = __ldg(args.coupling_entry + probe_begin + output);
+      const int components = __ldg(args.coupling_entry + nonzero_begin + term);
       const int indices[3] = {selection & 0xFF, (selection >> 8) & 0xFF,
                               (selection >> 16) & 0xFF};
-      const bool active[3] = {degrees[0] == degree && indices[0] == rank_index,
-                              degrees[1] == degree && indices[1] == rank_index,
-                              degrees[2] == degree && indices[2] == rank_index};
-      if (!active[0] && !active[1] && !active[2]) {
-        continue;
-      }
-      const float upstream =
+      const int component[3] = {components & 0xFF, (components >> 8) & 0xFF,
+                                (components >> 16) & 0xFF};
+      const float weight =
           __ldg(args.coupling_value + probe_begin + output) *
+          __ldg(args.coupling_value + nonzero_begin + term) *
           load_output_gradient(args.descriptor_gradient, args.output_inv_std,
                                node, P::OutputWidth, coordinate + output);
-      for (int term = 0; term < nonzero_count; ++term) {
-        const int components =
-            __ldg(args.coupling_entry + nonzero_begin + term);
-        const int component[3] = {components & 0xFF, (components >> 8) & 0xFF,
-                                  (components >> 16) & 0xFF};
-        const float weight =
-            upstream * __ldg(args.coupling_value + nonzero_begin + term);
-        const float first = probe_value(
-            probes, moments, degrees[0], component[0], indices[0]);
-        const float second = probe_value(
-            probes, moments, degrees[1], component[1], indices[1]);
-        const float third = probe_value(
-            probes, moments, degrees[2], component[2], indices[2]);
-        if (active[0]) {
-          scratch[component[0]] += weight * second * third;
-        }
-        if (active[1]) {
-          scratch[component[1]] += weight * first * third;
-        }
-        if (active[2]) {
-          scratch[component[2]] += weight * first * second;
-        }
+      float operand[3];
+#pragma unroll
+      for (int slot = 0; slot < 3; ++slot) {
+        operand[slot] = probe_value(
+            probes, moments, degrees[slot], component[slot], indices[slot]);
       }
+#pragma unroll
+      for (int slot = 0; slot < 3; ++slot) {
+        row[coupling_target(
+            degrees[slot], component[slot], indices[slot])] +=
+            weight * operand[(slot + 1) % 3] * operand[(slot + 2) % 3];
+      }
+      term += Width;
+      while (term >= nonzero_count) {
+        term -= nonzero_count;
+        ++output;
+      }
+    }
+  }
+  __syncthreads();
+  for (int target = lane; target < Targets; target += Width) {
+    float sum = 0.0f;
+    for (int source = 0; source < Width; ++source) {
+      sum += rows[source * Targets + target];
     }
+    total[target] = sum;
   }
+  __syncthreads();
 }
 
 // Four independent lane groups share one warp. An incomplete final block
@@ -919,7 +2403,6 @@ __global__ __launch_bounds__(
     Profile::Threads,
     2) void node_backward_kernel(Arguments args) {
   using P = Profile;
-  constexpr int MaxComponents = 9;
   const int thread = threadIdx.x;
   const int group = thread / P::NodeWidth;
   const int lane = thread & (P::NodeWidth - 1);
@@ -927,13 +2410,24 @@ __global__ __launch_bounds__(
   const bool active = candidate < args.node_count;
   const long node = active ? candidate : args.node_count - 1;
   const long state_offset = node * P::StateWidth;
+  if constexpr (Lmax >= 3) {
+    // The stages below read the descriptor cotangent of the node many times,
+    // mostly from loops whose loads feed the arithmetic of the same
+    // iteration, so the first touch of each of its lines would otherwise wait
+    // for the one before it. Degrees three and four hold a shared-memory
+    // footprint that leaves few warps per multiprocessor to hide those waits,
+    // so the node's lines are requested up front. The output scale and the
+    // readout matrices are shared by every node and stay cached after the
+    // first warp has read them.
+    prefetch_lines(
+        args.descriptor_gradient + node * P::OutputWidth, P::OutputWidth, lane,
+        P::NodeWidth);
+  }
 
   __shared__ float moments_storage[P::NodeGroups * P::MomentWidth];
   float* moments = moments_storage + group * P::MomentWidth;
-  for (int coordinate = lane; coordinate < P::MomentWidth;
-       coordinate += P::NodeWidth) {
-    moments[coordinate] = __ldg(args.state + state_offset + coordinate);
-  }
+  stage_shared(moments, args.state + state_offset,
+                                             lane);
   const float scalar_norm = __ldg(args.state + state_offset + P::MomentWidth);
   const float angular_norm =
       __ldg(args.state + state_offset + P::MomentWidth + 1);
@@ -999,17 +2493,20 @@ __global__ __launch_bounds__(
   }
   __syncthreads();
 
-  // Sparse-coupling scratch, one private row per lane.
+  // Sparse-coupling VJP of the node, one private accumulation row per lane
+  // and the folded total per node group.
+  constexpr int CouplingTargets = P::ProbeWidth + P::HighCount;
   __shared__ float
-      coupling_scratch[Lmax >= 3 ? P::NodeGroups * P::NodeWidth * MaxComponents
-                                 : 1];
-  float* scratch = nullptr;
+      coupling_rows[Lmax >= 3 ? P::NodeGroups * P::NodeWidth * CouplingTargets
+                              : 1];
+  __shared__ float
+      coupling_storage[Lmax >= 3 ? P::NodeGroups * CouplingTargets : 1];
+  const float* coupling = coupling_storage + group * CouplingTargets;
   if constexpr (Lmax >= 3) {
-    scratch = coupling_scratch + (group * P::NodeWidth + lane) * MaxComponents;
-#pragma unroll
-    for (int component = 0; component < MaxComponents; ++component) {
-      scratch[component] = 0.0f;
-    }
+    coupling_gradient(
+        lane, node, args, probes, moments,
+        coupling_rows + group * P::NodeWidth * CouplingTargets,
+        coupling_storage + group * CouplingTargets);
   }
 
   __shared__ float d_probes_storage[P::NodeGroups * P::ProbeWidth];
@@ -1074,12 +2571,9 @@ __global__ __launch_bounds__(
       }
     }
     if constexpr (Lmax >= 3) {
-      accumulate_coupling_gradient(
-          node, args, probes, moments, 1, lane, scratch);
 #pragma unroll
       for (int component = 0; component < 3; ++component) {
-        d_vector[component] += scratch[component];
-        scratch[component] = 0.0f;
+        d_vector[component] += coupling[component * P::K1 + lane];
       }
     }
 #pragma unroll
@@ -1154,12 +2648,9 @@ __global__ __launch_bounds__(
       }
     }
     if constexpr (Lmax >= 3) {
-      accumulate_coupling_gradient(
-          node, args, probes, moments, 2, lane, scratch);
 #pragma unroll
       for (int component = 0; component < 5; ++component) {
-        d_tensor[component] += scratch[component];
-        scratch[component] = 0.0f;
+        d_tensor[component] += coupling[3 * P::K1 + component * P::K2 + lane];
       }
     }
 #pragma unroll
@@ -1296,11 +2787,11 @@ __global__ __launch_bounds__(
           2.0f * load_output_gradient(args.descriptor_gradient,
                                       args.output_inv_std, node, P::OutputWidth,
                                       P::OutputGram3 + lane);
-      accumulate_coupling_gradient(
-          node, args, probes, moments, degree, 0, scratch);
+      const float* high_coupling =
+          coupling + P::ProbeWidth + (lane == 0 ? 0 : P::High3);
       for (int component = 0; component < count; ++component) {
         d_moments[offset + component] =
-            fmaf(gram, moments[offset + component], scratch[component]);
+            fmaf(gram, moments[offset + component], high_coupling[component]);
       }
     }
   }
@@ -1465,32 +2956,70 @@ __global__ __launch_bounds__(
 
 // === Edge recomputation backward ===
 
+/// Zero the stored cotangents of an edge that contributes nothing. A backward
+/// that assembles the force itself may skip either store.
+template 
+__device__ __forceinline__ void clear_edge_cotangents(const Arguments& args,
+                                                      long edge) {
+  if (args.edge_gradient != nullptr) {
+    args.edge_gradient[edge * 3 + 0] = 0.0f;
+    args.edge_gradient[edge * 3 + 1] = 0.0f;
+    args.edge_gradient[edge * 3 + 2] = 0.0f;
+  }
+  if constexpr (HasSpin) {
+    if (args.edge_spin_gradient != nullptr) {
+      args.edge_spin_gradient[edge * 3 + 0] = 0.0f;
+      args.edge_spin_gradient[edge * 3 + 1] = 0.0f;
+      args.edge_spin_gradient[edge * 3 + 2] = 0.0f;
+    }
+  }
+}
+
 template 
+          typename index_t,
+          int Warps,
+          int Split,
+          Scan Fold>
 __global__ __launch_bounds__(
-    Profile::Threads,
-    edge_resident_blocks()) void edge_backward_kernel(Arguments args) {
+    Warps * Split * kWarpSize,
+    (slots_types(Fold)
+         ? slot_resident_warps() ? Warps : 0>()
+         : edge_resident_blocks()) /
+        (Warps * Split)) void edge_backward_kernel(Arguments args) {
   using P = Profile;
+  enter_edge_scan(
+      args);
+  constexpr bool SlotScan = slots_types(Fold);
   constexpr int EdgeWidth = P::BackwardEdgeWidth;
   constexpr int Groups = kWarpSize / EdgeWidth;
   constexpr int ChannelTiles = Channels / EdgeWidth;
   constexpr int AngularTiles = (P::C1 + EdgeWidth - 1) / EdgeWidth;
   constexpr int TensorTiles = (P::C2 + EdgeWidth - 1) / EdgeWidth;
   constexpr int HighTiles = (P::HighCount + EdgeWidth - 1) / EdgeWidth;
+  static_assert(
+      P::Basis == 0 || (P::C1 % EdgeWidth == 0 && P::Basis % EdgeWidth == 0),
+      "the low-rank columns tile the lanes of an edge group");
 
-  const int thread = threadIdx.x;
-  const long node = blockIdx.x;
-  if (node >= args.node_count) {
-    return;
-  }
+  static_assert(Warps == 1 || Split == 1,
+                "a block either holds several rows or shares one");
+  const int thread = threadIdx.x % kWarpSize;
+  // A block holds `Warps` rows, one warp each, or one row shared by `Split`
+  // warps; `part` is the warp's share of a shared row.
+  const int warp = Warps > 1 ? threadIdx.x / kWarpSize : 0;
+  const int part = Split > 1 ? threadIdx.x / kWarpSize : 0;
+  constexpr bool Windowed = windowed();
+  constexpr int WalkWarps = Windowed ? Warps : 1;
+  const long node = static_cast(blockIdx.x) * Warps + warp;
+  // A warp past the last node of the final block owns no row: it reads node
+  // zero for its constants and scans an empty row, so that it still reaches
+  // every barrier of its block.
+  const bool active = node < args.node_count;
+  const long row_node = active ? node : 0;
   const int base_channel = thread & (EdgeWidth - 1);
   const int group = thread / EdgeWidth;
   const int leader = group * EdgeWidth;
@@ -1498,15 +3027,18 @@ __global__ __launch_bounds__(
   const auto* edge_index = static_cast(args.edge_index);
   const auto* destination_order =
       static_cast(args.destination_order);
-  const int center_type = static_cast(args.atype[args.node_begin + node]);
-  // The centre half of every pair length scale is a block constant, like the
-  // centre type it follows.
+  const int center_type =
+      static_cast(args.atype[args.node_begin + row_node]);
+  // The centre half of every pair length scale is a node constant, like the
+  // centre type it follows, and so is the length from which the window of no
+  // pair of the row needs resolving.
   const float center_radius =
       Bridged ? __ldg(args.contact_radius + center_type) : 0.0f;
-  const long begin = args.destination_row_ptr[node];
-  const long end = args.destination_row_ptr[node + 1];
+  const float bound = Bridged ? window_bound(args, center_radius) : 0.0f;
+  const long begin = args.destination_row_ptr[row_node];
+  const long end = active ? args.destination_row_ptr[row_node + 1] : begin;
   const int radial_modes = HasModes ? args.radial_modes : 0;
-  const long gradient_offset = node * P::StateWidth;
+  const long gradient_offset = row_node * P::StateWidth;
   const float scalar_mass_gradient =
       __ldg(args.moment_gradient + gradient_offset + P::MomentWidth);
   const float angular_mass_gradient =
@@ -1517,25 +3049,29 @@ __global__ __launch_bounds__(
   // cotangent fold into one weight of the radial slope.
   const bool has_pair = Bridged && args.pair_table != nullptr;
   const float pair_weight =
-      has_pair ? 0.5f * static_cast(args.pair_seed[node]) : 0.0f;
+      has_pair ? 0.5f * static_cast(args.pair_seed[row_node]) : 0.0f;
   float pair_energy = 0.0f;
 
   // The scalar cotangent is one vector of width C0 that every edge rereads.
   // Holding it in registers costs one entry per channel tile and spills the
   // widest profiles; shared memory serves it as a conflict-free broadcast
-  // because concurrent edge groups address identical channels.
-  __shared__ float scalar_gradient[Channels];
-  for (int channel = thread; channel < Channels; channel += P::Threads) {
+  // because concurrent edge groups address identical channels. The warps of
+  // a shared row stage it together, each its own share of the channels.
+  __shared__ float scalar_gradient_storage[Warps][Channels];
+  float* const scalar_gradient = scalar_gradient_storage[warp];
+  for (int channel = part * kWarpSize + thread; channel < Channels;
+       channel += Split * kWarpSize) {
     scalar_gradient[channel] = __ldg(args.moment_gradient + gradient_offset +
                                      P::ScalarOffset + channel);
   }
   // The reduced spin cotangent is a node constant that every edge rereads and
   // that every lane of a group needs in full, so it is staged in shared
   // memory for the same reason as the scalar cotangent above.
-  __shared__ float spin_gradient[HasSpin ? P::SpinEdgeWidth : 1];
+  __shared__ float spin_gradient_storage[Warps][HasSpin ? P::SpinEdgeWidth : 1];
+  float* const spin_gradient = spin_gradient_storage[warp];
   if constexpr (HasSpin) {
-    for (int coordinate = thread; coordinate < P::SpinEdgeWidth;
-         coordinate += P::Threads) {
+    for (int coordinate = part * kWarpSize + thread;
+         coordinate < P::SpinEdgeWidth; coordinate += Split * kWarpSize) {
       spin_gradient[coordinate] = __ldg(args.moment_gradient + gradient_offset +
                                         P::SpinOffset + coordinate);
     }
@@ -1544,50 +3080,336 @@ __global__ __launch_bounds__(
   float d_tensor[TensorTiles][5] = {};
   float d_high[HighTiles > 0 ? HighTiles : 1] = {};
 #pragma unroll
-  for (int tile = 0; tile < AngularTiles; ++tile) {
-    const int channel = base_channel + tile * EdgeWidth;
-    if (channel < P::C1) {
+  for (int tile = 0; tile < AngularTiles; ++tile) {
+    const int channel = base_channel + tile * EdgeWidth;
+    if (channel < P::C1) {
+#pragma unroll
+      for (int component = 0; component < 3; ++component) {
+        d_vector[tile][component] =
+            __ldg(args.moment_gradient + gradient_offset + P::VectorOffset +
+                  component * P::C1 + channel);
+      }
+    }
+  }
+#pragma unroll
+  for (int tile = 0; tile < TensorTiles; ++tile) {
+    const int channel = base_channel + tile * EdgeWidth;
+    if (channel < P::C2) {
+#pragma unroll
+      for (int component = 0; component < 5; ++component) {
+        d_tensor[tile][component] =
+            __ldg(args.moment_gradient + gradient_offset + P::TensorOffset +
+                  component * P::C2 + channel);
+      }
+    }
+  }
+  if constexpr (Lmax >= 3) {
+#pragma unroll
+    for (int component = 0; component < P::HighCount; ++component) {
+      if (component % EdgeWidth == base_channel) {
+        d_high[component / EdgeWidth] = __ldg(
+            args.moment_gradient + gradient_offset + P::HighOffset + component);
+      }
+    }
+  }
+
+  // Explicitly aligned: the mode residual reads these rows as float4, and
+  // a preceding shared array of arbitrary length would otherwise leave the
+  // block on a four-byte boundary.
+  __shared__ __align__(
+      16) float mode_cache[Warps * Split][HasModes ? Groups * kModeStride : 1];
+  __shared__ __align__(16) float
+      mode_derivative_cache[Warps * Split][HasModes ? Groups * kModeStride : 1];
+  const int mode_offset = HasModes ? group * kModeStride : 0;
+  float* modes = mode_cache[warp + part] + mode_offset;
+  float* mode_derivatives = mode_derivative_cache[warp + part] + mode_offset;
+  // A co-scheduled block of a slotted low-rank scan stages the transpose of
+  // its expansion matrix for the projections of the cotangents of its rows
+  // onto the basis functions of the slots. The blocks of the single-type
+  // scans, which also serve the systems that scan every edge pairwise and
+  // never project, and the smaller blocks keep the shared memory to the L1
+  // cache that their table reads use.
+  constexpr bool StageBasis = slots_basis(Fold) && windowed();
+  __shared__ float
+      transposed_basis[StageBasis ? (P::C0 - P::C1) * (P::Basis + 1) : 1];
+  if constexpr (StageBasis) {
+    stage_transposed_basis
(args.basis_matrix, transposed_basis);
+    __syncthreads();
+  } else {
+    // The warps of a shared row read the node cotangents once all of them
+    // have staged their shares.
+    sync_row();
+  }
+
+  // === Folded cotangents of a single-type system ===
+  // With the pair fixed, the coefficients move from the edges onto the node
+  // cotangents: the FiLM scale multiplies the channel cotangents, and each
+  // extra feature receives the coefficient-weighted sums of the channel
+  // cotangents, held by the lane that owns the feature in `extra`.
+  constexpr bool Foldable = foldable();
+  constexpr int ExtraFeatures = (HasModes ? kFoldedModes : 0) + 1;
+  constexpr int ExtraTiles = (ExtraFeatures + EdgeWidth - 1) / EdgeWidth;
+  float extra[ExtraTiles][9] = {};
+  int fold_type = -1;
+  if constexpr (Foldable && !SlotScan) {
+    fold_type = single_type(args, radial_modes);
+  }
+  const bool folded = fold_type >= 0;
+  // A slotted scan reads the cotangents of the mode profiles and of the basis
+  // functions of the slot of the source type of every edge.
+  static_assert(
+      !SlotScan || slottable(),
+      "a slotted scan serves a profile that slots");
+  constexpr bool SlotModes = SlotScan && HasModes;
+  constexpr bool SlotBasis = slots_basis(Fold);
+  constexpr int Slots = type_slots(Fold);
+  // A spin profile keeps the slot cotangents of its modes in shared memory,
+  // one row per mode of a slot, which the lanes of the modes of an edge group
+  // read at the slot of every edge: the spin families leave no registers for
+  // them. A row holds the nine cotangents of the harmonic components padded
+  // to twelve floats, so that two `float4` loads and one float read it. A
+  // spin-free profile keeps them in registers, in the slot tiles of
+  // `slot_extra`.
+  constexpr bool SharedSlots = SlotModes && HasSpin;
+  constexpr int CotangentStride = 12;
+  constexpr int SlotTiles =
+      SlotModes && !SharedSlots ? slot_tiles() : 1;
+  [[maybe_unused]] constexpr int SlotLanes = EdgeWidth / kFoldedModes;
+  static_assert(
+      !SlotModes || (EdgeWidth % kFoldedModes == 0 && Slots % SlotLanes == 0),
+      "the modes of the slots tile the lanes of an edge group");
+  [[maybe_unused]] float slot_extra[SlotTiles][9] = {};
+  __shared__ __align__(16) float slot_cotangent_storage
+      [Warps][SharedSlots ? Slots * kFoldedModes * CotangentStride : 1];
+  [[maybe_unused]] float* const slot_cotangent = slot_cotangent_storage[warp];
+  TypeSlots slots;
+  if constexpr (SlotScan) {
+    slots = load_type_slots(args);
+  }
+  // Cotangents of the basis functions of a low-rank single-type scan, which
+  // replace those of the channels they expand.
+  __shared__ float basis_gradient_storage[Warps]
+                                         [Fold == Scan::LowRank ? P::Basis : 1];
+  float* const basis_gradient = basis_gradient_storage[warp];
+  // Cotangents of the basis functions of each slot of a slotted low-rank
+  // scan, in rows of `Basis + 1` floats: the basis functions, then the
+  // cotangent of the FiLM shift of the channels from `C1` on of the slot. A
+  // lane holds `BasisTiles` basis functions.
+  constexpr int BasisTiles = SlotBasis ? P::Basis / EdgeWidth : 1;
+  constexpr int BasisStride = P::Basis + 1;
+  __shared__ float slot_basis_storage[Warps]
+                                     [SlotBasis ? Slots * BasisStride : 1];
+  float* const slot_basis_gradient = slot_basis_storage[warp];
+  float2 leading_film = make_float2(0.0f, 0.0f);
+  const float* leading_mixing = nullptr;
+  if (folded) {
+    const long pair =
+        static_cast(center_type) * args.type_count + fold_type;
+    const float2* film_row =
+        reinterpret_cast(args.pair_film + pair * Channels * 2) +
+        base_channel;
+    const float* mixing_row =
+        HasModes
+            ? args.pair_mixing + (pair * Channels + base_channel) * radial_modes
+            : nullptr;
+    leading_film = __ldg(reinterpret_cast(args.pair_film) +
+                         pair * Channels);
+    leading_mixing =
+        HasModes ? args.pair_mixing + pair * Channels * radial_modes : nullptr;
+#pragma unroll
+    for (int feature = 0; feature < ExtraFeatures; ++feature) {
+      if (feature <= radial_modes) {
+        float sums[9] = {};
+#pragma unroll
+        for (int tile = 0; tile < ChannelTiles; ++tile) {
+          const int channel = base_channel + tile * EdgeWidth;
+          const float weight =
+              feature == radial_modes
+                  ? __ldg(film_row + tile * EdgeWidth).y
+                  : __ldg(mixing_row + tile * EdgeWidth * radial_modes +
+                          feature);
+          sums[0] = fmaf(weight, scalar_gradient[channel], sums[0]);
+          if (channel < P::C1) {
+#pragma unroll
+            for (int component = 0; component < 3; ++component) {
+              sums[1 + component] =
+                  fmaf(weight, d_vector[tile][component], sums[1 + component]);
+            }
+          }
+          if (channel < P::C2) {
+#pragma unroll
+            for (int component = 0; component < 5; ++component) {
+              sums[4 + component] =
+                  fmaf(weight, d_tensor[tile][component], sums[4 + component]);
+            }
+          }
+        }
+#pragma unroll
+        for (int component = 0; component < 9; ++component) {
+          sums[component] = subwarp_allsum(sums[component], mask);
+          if (feature % EdgeWidth == base_channel) {
+            extra[feature / EdgeWidth][component] = sums[component];
+          }
+        }
+      }
+    }
+    // Every warp of the row has read the staged scalar cotangent before the
+    // first one scales it in place.
+    sync_row();
+#pragma unroll
+    for (int tile = 0; tile < ChannelTiles; ++tile) {
+      const int channel = base_channel + tile * EdgeWidth;
+      const float scale = __ldg(film_row + tile * EdgeWidth).x;
+      if (group == 0 && part == 0) {
+        scalar_gradient[channel] *= scale;
+      }
+      if (channel < P::C1) {
+#pragma unroll
+        for (int component = 0; component < 3; ++component) {
+          d_vector[tile][component] *= scale;
+        }
+      }
+      if (channel < P::C2) {
 #pragma unroll
-      for (int component = 0; component < 3; ++component) {
-        d_vector[tile][component] =
-            __ldg(args.moment_gradient + gradient_offset + P::VectorOffset +
-                  component * P::C1 + channel);
+        for (int component = 0; component < 5; ++component) {
+          d_tensor[tile][component] *= scale;
+        }
+      }
+    }
+    sync_row();
+    if constexpr (Fold == Scan::LowRank) {
+      if (part == 0) {
+        project_basis
(scalar_gradient, args.basis_matrix, basis_gradient,
+                         thread);
       }
+      sync_row();
     }
   }
+  if constexpr (SlotModes) {
+    // Each mode of a slot receives the sums of the channel cotangents
+    // weighted by the mode mixing of the pair of the slot, held by the lane
+    // that owns the mode of the slot in `slot_extra`.
 #pragma unroll
-  for (int tile = 0; tile < TensorTiles; ++tile) {
-    const int channel = base_channel + tile * EdgeWidth;
-    if (channel < P::C2) {
+    for (int slot = 0; slot < Slots; ++slot) {
+      if (slot < slots.count) {
+        const long pair = static_cast(center_type) * args.type_count +
+                          slot_type(slots, slot);
+        const float* mixing_row =
+            args.pair_mixing + (pair * Channels + base_channel) * radial_modes;
 #pragma unroll
-      for (int component = 0; component < 5; ++component) {
-        d_tensor[tile][component] =
-            __ldg(args.moment_gradient + gradient_offset + P::TensorOffset +
-                  component * P::C2 + channel);
+        for (int mode = 0; mode < kFoldedModes; ++mode) {
+          if (mode < radial_modes) {
+            float sums[9] = {};
+#pragma unroll
+            for (int tile = 0; tile < ChannelTiles; ++tile) {
+              const int channel = base_channel + tile * EdgeWidth;
+              const float weight =
+                  __ldg(mixing_row + tile * EdgeWidth * radial_modes + mode);
+              sums[0] = fmaf(weight, scalar_gradient[channel], sums[0]);
+              if (channel < P::C1) {
+#pragma unroll
+                for (int component = 0; component < 3; ++component) {
+                  sums[1 + component] = fmaf(weight, d_vector[tile][component],
+                                             sums[1 + component]);
+                }
+              }
+              if (channel < P::C2) {
+#pragma unroll
+                for (int component = 0; component < 5; ++component) {
+                  sums[4 + component] = fmaf(weight, d_tensor[tile][component],
+                                             sums[4 + component]);
+                }
+              }
+            }
+#pragma unroll
+            for (int component = 0; component < 9; ++component) {
+              sums[component] =
+                  subwarp_allsum(sums[component], mask);
+              if constexpr (SharedSlots) {
+                // Every warp of a shared row forms the same sums; the first
+                // one publishes them.
+                if (part == 0 && thread == component) {
+                  slot_cotangent[(slot * kFoldedModes + mode) *
+                                     CotangentStride +
+                                 component] = sums[component];
+                }
+              } else if (base_channel == slot_owner(slot, mode)) {
+                slot_extra[slot / SlotLanes][component] = sums[component];
+              }
+            }
+          }
+        }
       }
     }
+    if constexpr (SharedSlots) {
+      sync_row();
+    }
   }
-  if constexpr (Lmax >= 3) {
+  if constexpr (SlotBasis) {
+    // The cotangent of basis function `k` of a slot is the sum, over the
+    // channels from `C1` on in increasing order, of the channel cotangent
+    // times the FiLM scale of the pair of the slot times expansion
+    // coefficient `k` of the channel; that of the FiLM shift of the slot is
+    // the sum of the channel cotangents times the shifts. Lane `k` of the
+    // first warp of the row forms the cotangents of basis function `k` of
+    // every slot, and lane zero those of the shifts, for every warp of the
+    // row.
+    static_assert(P::Basis == kWarpSize,
+                  "every lane forms the cotangents of one basis function");
+    if (part == 0) {
+      constexpr int Expanded = P::C0 - P::C1;
+      const float2* film_rows[Slots];
 #pragma unroll
-    for (int component = 0; component < P::HighCount; ++component) {
-      if (component % EdgeWidth == base_channel) {
-        d_high[component / EdgeWidth] = __ldg(
-            args.moment_gradient + gradient_offset + P::HighOffset + component);
+      for (int slot = 0; slot < Slots; ++slot) {
+        const int type = slot_type(slots, slot < slots.count ? slot : 0);
+        film_rows[slot] =
+            reinterpret_cast(
+                args.pair_film +
+                (static_cast(center_type) * args.type_count + type) *
+                    Channels * 2) +
+            P::C1;
+      }
+      float basis_sums[Slots] = {};
+      float shift_sums[Slots] = {};
+      for (int first = 0; first < Expanded; first += 4) {
+        const float4 quad =
+            StageBasis ? staged_basis_quad
(transposed_basis, thread, first)
+                       : basis_quad
(args.basis_matrix, thread, first);
+        const float coefficients[4] = {quad.x, quad.y, quad.z, quad.w};
+#pragma unroll
+        for (int offset = 0; offset < 4; ++offset) {
+          const int channel = first + offset;
+          const float gradient = scalar_gradient[P::C1 + channel];
+#pragma unroll
+          for (int slot = 0; slot < Slots; ++slot) {
+            if (slot < slots.count) {
+              const float2 film = __ldg(film_rows[slot] + channel);
+              basis_sums[slot] = fmaf(coefficients[offset], gradient * film.x,
+                                      basis_sums[slot]);
+              shift_sums[slot] = fmaf(gradient, film.y, shift_sums[slot]);
+            }
+          }
+        }
+      }
+#pragma unroll
+      for (int slot = 0; slot < Slots; ++slot) {
+        slot_basis_gradient[slot * BasisStride + thread] = basis_sums[slot];
+        if (thread == 0) {
+          slot_basis_gradient[slot * BasisStride + P::Basis] = shift_sums[slot];
+        }
       }
     }
+    sync_row();
   }
 
-  // Explicitly aligned: the mode residual reads these rows as float4, and
-  // a preceding shared array of arbitrary length would otherwise leave the
-  // block on a four-byte boundary.
-  __shared__ __align__(
-      16) float mode_cache[HasModes ? Groups * kModeStride : 1];
-  __shared__ __align__(
-      16) float mode_derivative_cache[HasModes ? Groups * kModeStride : 1];
-  const int mode_offset = HasModes ? group * kModeStride : 0;
-  float* modes = mode_cache + mode_offset;
-  float* mode_derivatives = mode_derivative_cache + mode_offset;
-  __syncthreads();
+  // Fused assembly state. Each lane of an edge group owns the force and
+  // virial components congruent to its channel modulo the group width, so the
+  // edge totals are distributed rather than serialized on the leading lane.
+  constexpr int ForceSlots = (3 + EdgeWidth - 1) / EdgeWidth;
+  constexpr int VirialSlots = (9 + EdgeWidth - 1) / EdgeWidth;
+  const bool assemble = args.force != nullptr;
+  float destination_force[ForceSlots] = {};
+  float node_virial[VirialSlots] = {};
 
   // Offsets of the five spin families inside the staged cotangent row.
   constexpr int SpinMagnitudeSlot = P::SpinMagnitude - P::SpinOffset;
@@ -1596,20 +3418,24 @@ __global__ __launch_bounds__(
   constexpr int SpinBondSlot = P::SpinBond - P::SpinOffset;
   constexpr int SpinTensorSlot = P::SpinTensor - P::SpinOffset;
 
-  for (long position = begin + group; position < end; position += Groups) {
+  const auto scan_edge = [&](const long position, auto scan) {
+    constexpr Scan Kind = decltype(scan)::value;
+    constexpr bool Folded = Kind == Scan::Direct || Kind == Scan::LowRank;
+    constexpr int FoldColumns = fold_columns();
+    constexpr int FoldTiles = FoldColumns / EdgeWidth;
+    // The tiles of a low-rank scan from ``C1 / EdgeWidth`` on hold basis
+    // functions, whose cotangents replace those of the channels.
+    constexpr int DirectTiles =
+        Kind == Scan::LowRank ? P::C1 / EdgeWidth : FoldTiles;
+    // Channel tiles whose FiLM amplitude the edges carry.
+    constexpr int EdgeTiles =
+        slots_basis(Kind) ? P::C1 / EdgeWidth : ChannelTiles;
     const long edge = edge_at_position(position, destination_order);
     if (args.edge_mask != nullptr && !args.edge_mask[edge]) {
       if (thread == leader) {
-        args.edge_gradient[edge * 3 + 0] = 0.0f;
-        args.edge_gradient[edge * 3 + 1] = 0.0f;
-        args.edge_gradient[edge * 3 + 2] = 0.0f;
-        if constexpr (HasSpin) {
-          args.edge_spin_gradient[edge * 3 + 0] = 0.0f;
-          args.edge_spin_gradient[edge * 3 + 1] = 0.0f;
-          args.edge_spin_gradient[edge * 3 + 2] = 0.0f;
-        }
+        clear_edge_cotangents(args, edge);
       }
-      continue;
+      return;
     }
     const long source = static_cast(edge_index[edge]);
     const int source_type = static_cast(args.atype[source]);
@@ -1617,22 +3443,15 @@ __global__ __launch_bounds__(
       if (center_type >= args.type_count - 1 ||
           source_type >= args.type_count - 1) {
         if (thread == leader) {
-          args.edge_gradient[edge * 3 + 0] = 0.0f;
-          args.edge_gradient[edge * 3 + 1] = 0.0f;
-          args.edge_gradient[edge * 3 + 2] = 0.0f;
-          if constexpr (HasSpin) {
-            args.edge_spin_gradient[edge * 3 + 0] = 0.0f;
-            args.edge_spin_gradient[edge * 3 + 1] = 0.0f;
-            args.edge_spin_gradient[edge * 3 + 2] = 0.0f;
-          }
+          clear_edge_cotangents(args, edge);
         }
-        continue;
+        return;
       }
     }
     const EdgeGeometry geometry = load_geometry(
         edge, source, source_type, args.rcut, args.f_inner, args.f_outer,
-        args.inverse_f_width, center_radius, args.contact_radius, args.eps,
-        args.edge_vec);
+        args.inverse_f_width, center_radius, bound, args.contact_radius,
+        args.eps, args.edge_vec);
     const TableLocation location = locate_table(
         geometry.seen, args.table_stride, args.table_max, args.interval_count);
     // The pair term depends on the geometry alone, so it starts here and runs
@@ -1648,18 +3467,23 @@ __global__ __launch_bounds__(
       pair_energy = fmaf(0.5f, term.energy, pair_energy);
       pair_slope = term.slope;
     }
-    const TableRow row = table_row(args.table, location, args.table_width);
+    const TableRow row =
+        reads_low_rank(Kind)
+            ? table_row(args.basis_table, location, FoldColumns + radial_modes)
+            : table_row(args.table, location, args.table_width);
     const float coordinate = location.coordinate;
     const bool clamped = location.clamped;
     if constexpr (HasModes) {
-      __syncwarp(mask);
-      for (int mode = base_channel; mode < radial_modes; mode += EdgeWidth) {
-        const float2 value = evaluate_table_with_derivative(
-            row, Channels + mode, coordinate, clamped);
-        modes[mode] = value.x;
-        mode_derivatives[mode] = value.y;
+      if constexpr (Kind == Scan::Pairwise) {
+        __syncwarp(mask);
+        for (int mode = base_channel; mode < radial_modes; mode += EdgeWidth) {
+          const float2 value = evaluate_table_with_derivative(
+              row, Channels + mode, coordinate, clamped);
+          modes[mode] = value.x;
+          mode_derivatives[mode] = value.y;
+        }
+        __syncwarp(mask);
       }
-      __syncwarp(mask);
     }
 
     float basis[9];
@@ -1667,17 +3491,10 @@ __global__ __launch_bounds__(
     const float envelope = geometry.envelope;
     const long pair =
         static_cast(center_type) * args.type_count + geometry.source_type;
-    const float2* film_row =
-        reinterpret_cast(args.pair_film + pair * Channels * 2) +
-        base_channel;
-    const float* mixing_row =
-        HasModes
-            ? args.pair_mixing + (pair * Channels + base_channel) * radial_modes
-            : nullptr;
 
     // The high-degree cotangent enters the leading channel only, so its
     // angular contraction is reduced onto the lane that owns that channel.
-    float high_angular = 0.0f;
+    [[maybe_unused]] float high_angular = 0.0f;
     if constexpr (Lmax >= 3) {
 #pragma unroll
       for (int component = 0; component < P::HighCount; ++component) {
@@ -1693,13 +3510,13 @@ __global__ __launch_bounds__(
     // The spin payload of this edge, hoisted like the forward one: every
     // channel multiplies the same conditioned neighbour moment and the same
     // projection of that moment onto the bond.
-    float neighbor_spin[3] = {0.0f, 0.0f, 0.0f};
-    float neighbor_magnitude = 0.0f;
-    float neighbor_gate = 0.0f;
-    float neighbor_scale = 0.0f;
-    float bond_alignment = 0.0f;
-    float spin_cotangent[3] = {0.0f, 0.0f, 0.0f};
-    const float2* spin_row = nullptr;
+    [[maybe_unused]] float neighbor_spin[3] = {0.0f, 0.0f, 0.0f};
+    [[maybe_unused]] float neighbor_magnitude = 0.0f;
+    [[maybe_unused]] float neighbor_gate = 0.0f;
+    [[maybe_unused]] float neighbor_scale = 0.0f;
+    [[maybe_unused]] float bond_alignment = 0.0f;
+    [[maybe_unused]] float spin_cotangent[3] = {0.0f, 0.0f, 0.0f};
+    [[maybe_unused]] const float2* spin_row = nullptr;
     if constexpr (HasSpin) {
       const SpinTypeWeights weights =
           load_spin_type(args.spin_type, geometry.source_type);
@@ -1727,171 +3544,420 @@ __global__ __launch_bounds__(
     // reach the coordinate gradient through the one transverse projection in
     // `basis_vjp`, so no family is reduced separately.
     float direction_gradient[3] = {0.0f, 0.0f, 0.0f};
-    float angular_zero = 0.0f;
+    // Angular payload of the leading channel, held by the leader of the group,
+    // and the slope of that channel that the high degrees read: of its bare
+    // radial feature in a folded scan, and in the other scans of its FiLM
+    // amplitude before the mode mixing, whose value `leading_value` holds.
+    [[maybe_unused]] float leading_payload = 0.0f;
+    [[maybe_unused]] float leading_value = 0.0f;
+    [[maybe_unused]] float leading_slope = 0.0f;
+    if constexpr (Folded) {
+      // The pair coefficients live in the folded cotangents, so a channel
+      // contributes through its bare radial feature, and every extra feature
+      // through its own value on the lane that owns it. The angular cotangent
+      // accumulates without the squared envelope, which multiplies it once.
+      const float squared = envelope * envelope;
+      float leading = 0.0f;
 #pragma unroll
-    for (int tile = 0; tile < ChannelTiles; ++tile) {
-      const int channel = base_channel + tile * EdgeWidth;
-      const float2 radial =
-          evaluate_table_with_derivative(row, channel, coordinate, clamped);
-      const float2 film = __ldg(film_row + tile * EdgeWidth);
-      // ``film_value`` is the pre-envelope FiLM amplitude phi; the reduced
-      // payload is phi * chi for degree zero and phi * chi^2 above it.
-      float film_value = fmaf(film.x, radial.x, film.y);
-      float film_derivative = film.x * radial.y;
-      if constexpr (HasModes) {
-        accumulate_modes_with_derivative(
-            mixing_row + tile * (EdgeWidth * radial_modes), modes,
-            mode_derivatives, radial_modes, film_value, film_derivative);
-      }
-      float angular = 0.0f;
-      if (channel < P::C1) {
+      for (int tile = 0; tile < FoldTiles; ++tile) {
+        const int channel = base_channel + tile * EdgeWidth;
+        const float2 radial =
+            evaluate_table_with_derivative(row, channel, coordinate, clamped);
+        float angular = 0.0f;
+        if (channel < P::C1) {
 #pragma unroll
-        for (int component = 0; component < 3; ++component) {
-          angular =
-              fmaf(d_vector[tile][component], basis[1 + component], angular);
+          for (int component = 0; component < 3; ++component) {
+            angular =
+                fmaf(d_vector[tile][component], basis[1 + component], angular);
+          }
+        }
+        if (channel < P::C2) {
+#pragma unroll
+          for (int component = 0; component < 5; ++component) {
+            angular =
+                fmaf(d_tensor[tile][component], basis[4 + component], angular);
+          }
+        }
+        const float scaled = envelope * angular;
+        const float gradient = tile < DirectTiles
+                                   ? scalar_gradient[channel]
+                                   : basis_gradient[channel - P::C1];
+        radial_gradient =
+            fmaf(envelope * (gradient + scaled), radial.y, radial_gradient);
+        envelope_gradient =
+            fmaf(fmaf(2.0f, scaled, gradient), radial.x, envelope_gradient);
+        if (channel < P::C1) {
+#pragma unroll
+          for (int component = 0; component < 3; ++component) {
+            d_basis[1 + component] = fmaf(d_vector[tile][component], radial.x,
+                                          d_basis[1 + component]);
+          }
+        }
+        if (channel < P::C2) {
+#pragma unroll
+          for (int component = 0; component < 5; ++component) {
+            d_basis[4 + component] = fmaf(d_tensor[tile][component], radial.x,
+                                          d_basis[4 + component]);
+          }
+        }
+        if (tile == 0) {
+          leading = radial.x;
+          leading_slope = radial.y;
         }
       }
-      if (channel < P::C2) {
+      float feature_value[ExtraTiles] = {};
+      float feature_slope[ExtraTiles] = {};
 #pragma unroll
-        for (int component = 0; component < 5; ++component) {
-          angular =
-              fmaf(d_tensor[tile][component], basis[4 + component], angular);
+      for (int tile = 0; tile < ExtraTiles; ++tile) {
+        const int feature = base_channel + tile * EdgeWidth;
+        if (tile * EdgeWidth <= radial_modes) {
+          const float2 value =
+              feature < radial_modes
+                  ? evaluate_table_with_derivative(row, FoldColumns + feature,
+                                                   coordinate, clamped)
+                  : make_float2(feature == radial_modes ? 1.0f : 0.0f, 0.0f);
+          feature_value[tile] = value.x;
+          feature_slope[tile] = value.y;
+          float angular = 0.0f;
+#pragma unroll
+          for (int component = 1; component < 9; ++component) {
+            angular = fmaf(extra[tile][component], basis[component], angular);
+          }
+          const float scaled = envelope * angular;
+          radial_gradient = fmaf(envelope * (extra[tile][0] + scaled), value.y,
+                                 radial_gradient);
+          envelope_gradient = fmaf(fmaf(2.0f, scaled, extra[tile][0]), value.x,
+                                   envelope_gradient);
+#pragma unroll
+          for (int component = 1; component < 9; ++component) {
+            d_basis[component] =
+                fmaf(extra[tile][component], value.x, d_basis[component]);
+          }
         }
       }
       if constexpr (Lmax >= 3) {
-        if (channel == 0) {
-          angular += high_angular;
+        // The high-degree cotangent reaches the leading channel through its
+        // full amplitude, which the leader reassembles from the bare feature,
+        // its coefficients and the mode profiles of the group.
+        float film_value = fmaf(leading_film.x, leading, leading_film.y);
+        float film_slope = leading_film.x * leading_slope;
+        if constexpr (HasModes) {
+#pragma unroll
+          for (int mode = 0; mode < kFoldedModes; ++mode) {
+            if (mode < radial_modes) {
+              const int owner = leader + mode % EdgeWidth;
+              const float weight = __ldg(leading_mixing + mode);
+              film_value = fmaf(
+                  weight,
+                  __shfl_sync(mask, feature_value[mode / EdgeWidth], owner),
+                  film_value);
+              film_slope = fmaf(
+                  weight,
+                  __shfl_sync(mask, feature_slope[mode / EdgeWidth], owner),
+                  film_slope);
+            }
+          }
         }
+        if (thread == leader) {
+          radial_gradient =
+              fmaf(squared * high_angular, film_slope, radial_gradient);
+          envelope_gradient = fmaf(2.0f * envelope * high_angular, film_value,
+                                   envelope_gradient);
+        }
+        leading_payload = film_value * squared;
       }
-      // d/dphi   = chi * (p + chi * a)
-      // d/dchi   = phi * (p + 2 * chi * a)
-      const float scaled = envelope * angular;
-      const float film_gradient = scalar_gradient[channel] + scaled;
-      radial_gradient =
-          fmaf(envelope * film_gradient, film_derivative, radial_gradient);
-      envelope_gradient =
-          fmaf(film_gradient + scaled, film_value, envelope_gradient);
-      const float angular_payload = film_value * envelope * envelope;
-      if (channel < P::C1) {
 #pragma unroll
-        for (int component = 0; component < 3; ++component) {
-          d_basis[1 + component] =
-              fmaf(d_vector[tile][component], angular_payload,
-                   d_basis[1 + component]);
-        }
+      for (int component = 1; component < 9; ++component) {
+        d_basis[component] *= squared;
       }
-      if (channel < P::C2) {
+    } else {
+      const float2* film_row = reinterpret_cast(
+                                   args.pair_film + pair * Channels * 2) +
+                               base_channel;
+      [[maybe_unused]] const float* mixing_row =
+          HasModes ? args.pair_mixing +
+                         (pair * Channels + base_channel) * radial_modes
+                   : nullptr;
 #pragma unroll
-        for (int component = 0; component < 5; ++component) {
-          d_basis[4 + component] =
-              fmaf(d_tensor[tile][component], angular_payload,
-                   d_basis[4 + component]);
+      for (int tile = 0; tile < EdgeTiles; ++tile) {
+        const int channel = base_channel + tile * EdgeWidth;
+        const float2 radial =
+            evaluate_table_with_derivative(row, channel, coordinate, clamped);
+        const float2 film = __ldg(film_row + tile * EdgeWidth);
+        // ``film_value`` is the pre-envelope FiLM amplitude phi; the reduced
+        // payload is phi * chi for degree zero and phi * chi^2 above it.
+        float film_value = fmaf(film.x, radial.x, film.y);
+        float film_derivative = film.x * radial.y;
+        if constexpr (HasModes && Kind == Scan::Pairwise) {
+          accumulate_modes_with_derivative(
+              mixing_row + tile * (EdgeWidth * radial_modes), modes,
+              mode_derivatives, radial_modes, film_value, film_derivative);
         }
-      }
-      if constexpr (HasSpin) {
-        if (channel < P::Cs) {
-          const float2 spin_film = __ldg(spin_row + tile * EdgeWidth);
-          const float profile = fmaf(spin_film.x, radial.x, spin_film.y);
-          const float envelope_squared = envelope * envelope;
-          const float weight = profile * envelope_squared;
-          const float magnitude_gradient =
-              spin_gradient[SpinMagnitudeSlot + channel];
-          // The two degree-one cotangents are read once and used by the
-          // amplitude, the magnetic and the angular accumulations below.
-          float vector_gradient[3];
-          float bond_gradient[3];
+        float angular = 0.0f;
+        if (channel < P::C1) {
 #pragma unroll
           for (int component = 0; component < 3; ++component) {
-            vector_gradient[component] =
-                spin_gradient[SpinVectorSlot + component * P::Cs + channel];
-            bond_gradient[component] =
-                spin_gradient[SpinBondSlot + component * P::Cs + channel];
+            angular =
+                fmaf(d_vector[tile][component], basis[1 + component], angular);
           }
-          // Component of the bond cotangent along the bond. The bond-projected
-          // family is `P = w (s . u) u`, so this one dot product carries all
-          // three of its VJPs: `dP/dw` contracts to `(s . u) (dP . u)`,
-          // `dP/ds` to `w (dP . u) u`, and the direction term below reuses it
-          // once more.
-          float bond_projection = 0.0f;
+        }
+        if (channel < P::C2) {
 #pragma unroll
-          for (int component = 0; component < 3; ++component) {
-            bond_projection = fmaf(bond_gradient[component],
-                                   basis[1 + component], bond_projection);
+          for (int component = 0; component < 5; ++component) {
+            angular =
+                fmaf(d_tensor[tile][component], basis[4 + component], angular);
           }
-          float sigma = fmaf(
-              magnitude_gradient, neighbor_magnitude,
-              spin_gradient[SpinCoordinationSlot + channel] * neighbor_gate);
-          sigma = fmaf(bond_alignment, bond_projection, sigma);
-#pragma unroll
-          for (int component = 0; component < 3; ++component) {
-            sigma = fmaf(vector_gradient[component], neighbor_spin[component],
-                         sigma);
+        }
+        if constexpr (Lmax >= 3 && Kind == Scan::Pairwise) {
+          if (channel == 0) {
+            angular += high_angular;
           }
-          // Magnetic cotangent of the source moment. The coordination family
-          // reads the type gate rather than the moment, so it contributes
-          // nothing here.
-          const float linear = 2.0f * weight * magnitude_gradient;
-          const float projected = weight * bond_projection;
+        }
+        // d/dphi   = chi * (p + chi * a)
+        // d/dchi   = phi * (p + 2 * chi * a)
+        const float scaled = envelope * angular;
+        const float film_gradient = scalar_gradient[channel] + scaled;
+        radial_gradient =
+            fmaf(envelope * film_gradient, film_derivative, radial_gradient);
+        envelope_gradient =
+            fmaf(film_gradient + scaled, film_value, envelope_gradient);
+        const float angular_payload = film_value * envelope * envelope;
+        if (channel < P::C1) {
 #pragma unroll
           for (int component = 0; component < 3; ++component) {
-            spin_cotangent[component] =
-                fmaf(linear, neighbor_spin[component],
-                     fmaf(projected, basis[1 + component],
-                          fmaf(weight, vector_gradient[component],
-                               spin_cotangent[component])));
+            d_basis[1 + component] =
+                fmaf(d_vector[tile][component], angular_payload,
+                     d_basis[1 + component]);
           }
-          // Angular cotangent of the bond-projected family. Differentiating
-          // `(s . u) u` with respect to `u` at fixed `s` gives the rank-one
-          // sum `s (dP . u) + (s . u) dP`, which the transverse projection in
-          // `basis_vjp` turns into the coordinate gradient. Every other spin
-          // family is independent of the direction and contributes nothing.
-          const float aligned_weight = weight * bond_alignment;
+        }
+        if (channel < P::C2) {
 #pragma unroll
-          for (int component = 0; component < 3; ++component) {
-            direction_gradient[component] =
-                fmaf(projected, neighbor_spin[component],
-                     fmaf(aligned_weight, bond_gradient[component],
-                          direction_gradient[component]));
+          for (int component = 0; component < 5; ++component) {
+            d_basis[4 + component] =
+                fmaf(d_tensor[tile][component], angular_payload,
+                     d_basis[4 + component]);
           }
-          if (channel == 0) {
-            // The single neighbour quadrupole rides the leading spin channel.
-            // Both of its contributions are formed here rather than hoisted,
-            // so its five cotangent components never stay live across the
-            // channel loop.
+        }
+        if constexpr (HasSpin) {
+          if (channel < P::Cs) {
+            const float2 spin_film = __ldg(spin_row + tile * EdgeWidth);
+            const float profile = fmaf(spin_film.x, radial.x, spin_film.y);
+            const float envelope_squared = envelope * envelope;
+            const float weight = profile * envelope_squared;
+            const float magnitude_gradient =
+                spin_gradient[SpinMagnitudeSlot + channel];
+            // The two degree-one cotangents are read once and used by the
+            // amplitude, the magnetic and the angular accumulations below.
+            float vector_gradient[3];
+            float bond_gradient[3];
 #pragma unroll
-            for (int component = 0; component < 5; ++component) {
-              sigma = fmaf(
-                  spin_gradient[SpinTensorSlot + component],
-                  spin_quadrupole_component(neighbor_spin[0], neighbor_spin[1],
-                                            neighbor_spin[2], component),
-                  sigma);
+            for (int component = 0; component < 3; ++component) {
+              vector_gradient[component] =
+                  spin_gradient[SpinVectorSlot + component * P::Cs + channel];
+              bond_gradient[component] =
+                  spin_gradient[SpinBondSlot + component * P::Cs + channel];
+            }
+            // Component of the bond cotangent along the bond. The
+            // bond-projected family is `P = w (s . u) u`, so this one dot
+            // product carries all three of its VJPs: `dP/dw` contracts to
+            // `(s . u) (dP . u)`, `dP/ds` to `w (dP . u) u`, and the direction
+            // term below reuses it once more.
+            float bond_projection = 0.0f;
+#pragma unroll
+            for (int component = 0; component < 3; ++component) {
+              bond_projection = fmaf(bond_gradient[component],
+                                     basis[1 + component], bond_projection);
+            }
+            float sigma = fmaf(
+                magnitude_gradient, neighbor_magnitude,
+                spin_gradient[SpinCoordinationSlot + channel] * neighbor_gate);
+            sigma = fmaf(bond_alignment, bond_projection, sigma);
+#pragma unroll
+            for (int component = 0; component < 3; ++component) {
+              sigma = fmaf(vector_gradient[component], neighbor_spin[component],
+                           sigma);
+            }
+            // Magnetic cotangent of the source moment. The coordination family
+            // reads the type gate rather than the moment, so it contributes
+            // nothing here.
+            const float linear = 2.0f * weight * magnitude_gradient;
+            const float projected = weight * bond_projection;
+#pragma unroll
+            for (int component = 0; component < 3; ++component) {
+              spin_cotangent[component] =
+                  fmaf(linear, neighbor_spin[component],
+                       fmaf(projected, basis[1 + component],
+                            fmaf(weight, vector_gradient[component],
+                                 spin_cotangent[component])));
             }
-            float quadrupole[3];
-            spin_quadrupole_vjp(spin_gradient + SpinTensorSlot, weight,
-                                neighbor_spin, quadrupole);
+            // Angular cotangent of the bond-projected family. Differentiating
+            // `(s . u) u` with respect to `u` at fixed `s` gives the rank-one
+            // sum `s (dP . u) + (s . u) dP`, which the transverse projection in
+            // `basis_vjp` turns into the coordinate gradient. Every other spin
+            // family is independent of the direction and contributes nothing.
+            const float aligned_weight = weight * bond_alignment;
 #pragma unroll
             for (int component = 0; component < 3; ++component) {
-              spin_cotangent[component] += quadrupole[component];
+              direction_gradient[component] =
+                  fmaf(projected, neighbor_spin[component],
+                       fmaf(aligned_weight, bond_gradient[component],
+                            direction_gradient[component]));
             }
+            if (channel == 0) {
+              // The single neighbour quadrupole rides the leading spin channel.
+              // Both of its contributions are formed here rather than hoisted,
+              // so its five cotangent components never stay live across the
+              // channel loop.
+#pragma unroll
+              for (int component = 0; component < 5; ++component) {
+                sigma = fmaf(spin_gradient[SpinTensorSlot + component],
+                             spin_quadrupole_component(
+                                 neighbor_spin[0], neighbor_spin[1],
+                                 neighbor_spin[2], component),
+                             sigma);
+              }
+              float quadrupole[3];
+              spin_quadrupole_vjp(spin_gradient + SpinTensorSlot, weight,
+                                  neighbor_spin, quadrupole);
+#pragma unroll
+              for (int component = 0; component < 3; ++component) {
+                spin_cotangent[component] += quadrupole[component];
+              }
+            }
+            // w = chi^2 (gamma g + beta): the distance enters through the table
+            // and through the envelope, and `sigma` is the cotangent of that
+            // shared amplitude over all five families.
+            radial_gradient = fmaf(envelope_squared * sigma,
+                                   spin_film.x * radial.y, radial_gradient);
+            envelope_gradient =
+                fmaf(2.0f * envelope * profile, sigma, envelope_gradient);
+          }
+        }
+        if constexpr (Lmax >= 3) {
+          if (channel == 0) {
+            leading_payload = angular_payload;
+            leading_value = film_value;
+            leading_slope = film_derivative;
           }
-          // w = chi^2 (gamma g + beta): the distance enters through the table
-          // and through the envelope, and `sigma` is the cotangent of that
-          // shared amplitude over all five families.
-          radial_gradient = fmaf(envelope_squared * sigma,
-                                 spin_film.x * radial.y, radial_gradient);
-          envelope_gradient =
-              fmaf(2.0f * envelope * profile, sigma, envelope_gradient);
         }
       }
-      if constexpr (Lmax >= 3) {
-        if (channel == 0) {
-          angular_zero = angular_payload;
+      if constexpr (slots_types(Kind)) {
+        const int slot = type_slot(slots, geometry.source_type);
+        [[maybe_unused]] const float squared = envelope * envelope;
+        [[maybe_unused]] float2 profile = make_float2(0.0f, 0.0f);
+        if constexpr (HasModes) {
+          // The mode of the lane contributes through its bare profile, on the
+          // lanes that own the slot of the source type, as an extra feature
+          // of a single-type scan does.
+          const int mode = base_channel % kFoldedModes;
+          if (mode < radial_modes) {
+            profile = evaluate_table_with_derivative(row, FoldColumns + mode,
+                                                     coordinate, clamped);
+          }
+          // The contribution of mode `mode` of the slot of the edge, from its
+          // nine cotangents `cotangent`.
+          const auto add_mode = [&](const float (&cotangent)[9]) {
+            float angular = 0.0f;
+#pragma unroll
+            for (int component = 1; component < 9; ++component) {
+              angular = fmaf(cotangent[component], basis[component], angular);
+            }
+            const float scaled = envelope * angular;
+            radial_gradient = fmaf(envelope * (cotangent[0] + scaled),
+                                   profile.y, radial_gradient);
+            envelope_gradient = fmaf(fmaf(2.0f, scaled, cotangent[0]),
+                                     profile.x, envelope_gradient);
+            const float weight = profile.x * squared;
+#pragma unroll
+            for (int component = 1; component < 9; ++component) {
+              d_basis[component] =
+                  fmaf(cotangent[component], weight, d_basis[component]);
+            }
+          };
+          if constexpr (SharedSlots) {
+            // The lanes of the modes of the descriptor read the row of their
+            // mode of the slot of the edge; the rows of the modes past
+            // `radial_modes` hold no cotangents. A slotted launch carries at
+            // most `kFoldedModes` modes, so lane `l < radial_modes` is the
+            // lane of mode `l`.
+            if (base_channel < radial_modes) {
+              const float* const cotangent_row =
+                  slot_cotangent +
+                  (slot * kFoldedModes + mode) * CotangentStride;
+              const float4 low =
+                  reinterpret_cast(cotangent_row)[0];
+              const float4 high =
+                  reinterpret_cast(cotangent_row)[1];
+              const float cotangent[9] = {low.x,  low.y,  low.z,
+                                          low.w,  high.x, high.y,
+                                          high.z, high.w, cotangent_row[8]};
+              add_mode(cotangent);
+            }
+          } else {
+#pragma unroll
+            for (int tile = 0; tile < SlotTiles; ++tile) {
+              if (owned_slot(base_channel, tile) == slot) {
+                add_mode(slot_extra[tile]);
+              }
+            }
+          }
+        }
+        if constexpr (slots_basis(Kind)) {
+          // The basis functions of the lane contribute through their bare
+          // values with the cotangents of the slot of the source type, and
+          // the FiLM shift of the slot through the envelope on the leader.
+          const float* const basis_cotangent =
+              slot_basis_gradient + slot * BasisStride;
+#pragma unroll
+          for (int tile = 0; tile < BasisTiles; ++tile) {
+            const int function = base_channel + tile * EdgeWidth;
+            const float2 radial = evaluate_table_with_derivative(
+                row, P::C1 + function, coordinate, clamped);
+            const float gradient = basis_cotangent[function];
+            radial_gradient =
+                fmaf(envelope * gradient, radial.y, radial_gradient);
+            envelope_gradient = fmaf(gradient, radial.x, envelope_gradient);
+          }
+          if (thread == leader) {
+            envelope_gradient += basis_cotangent[P::Basis];
+          }
+        }
+        if constexpr (Lmax >= 3) {
+          // The high-degree cotangent reaches the leading channel through its
+          // full amplitude, which the leader completes with the mode mixing
+          // of the pair; lane `leader + r` of the group holds the profile of
+          // mode `r`.
+          float value = leading_value;
+          float slope = leading_slope;
+          if constexpr (HasModes) {
+            const float* const pair_modes =
+                args.pair_mixing + pair * Channels * radial_modes;
+#pragma unroll
+            for (int mixed = 0; mixed < kFoldedModes; ++mixed) {
+              if (mixed < radial_modes) {
+                const float weight = __ldg(pair_modes + mixed);
+                value =
+                    fmaf(weight, __shfl_sync(mask, profile.x, leader + mixed),
+                         value);
+                slope =
+                    fmaf(weight, __shfl_sync(mask, profile.y, leader + mixed),
+                         slope);
+              }
+            }
+          }
+          if (thread == leader) {
+            radial_gradient =
+                fmaf(squared * high_angular, slope, radial_gradient);
+            envelope_gradient =
+                fmaf(2.0f * envelope * high_angular, value, envelope_gradient);
+          }
+          leading_payload = value * squared;
         }
       }
     }
 
     if constexpr (Lmax >= 3) {
-      const float amplitude = __shfl_sync(mask, angular_zero, leader);
+      const float amplitude = __shfl_sync(mask, leading_payload, leader);
 #pragma unroll
       for (int component = 0; component < P::HighCount; ++component) {
         if (component % EdgeWidth == base_channel) {
@@ -1912,8 +3978,9 @@ __global__ __launch_bounds__(
     }
     // The table terms above differentiate the length the table reads; the
     // pair potential acts on the true length and joins after the chain rule.
-    radial_gradient = radial_chain(geometry, args, center_radius,
-                                            radial_gradient, envelope_gradient);
+    radial_gradient =
+        radial_chain(geometry, args, center_radius, bound,
+                              radial_gradient, envelope_gradient);
     if constexpr (Bridged) {
       radial_gradient = fmaf(pair_weight, pair_slope, radial_gradient);
     }
@@ -1925,33 +3992,161 @@ __global__ __launch_bounds__(
     basis_vjp(geometry, d_basis, direction_gradient, radial_gradient, output);
 #pragma unroll
     for (int component = 0; component < 3; ++component) {
-      output[component] = subwarp_sum(output[component], mask);
+      output[component] = subwarp_allsum(output[component], mask);
     }
-    if (thread == leader) {
+    if (thread == leader && args.edge_gradient != nullptr) {
       args.edge_gradient[edge * 3 + 0] = output[0];
       args.edge_gradient[edge * 3 + 1] = output[1];
       args.edge_gradient[edge * 3 + 2] = output[2];
     }
+    if (assemble) {
+      // The source half of the force belongs to a node this warp does not
+      // own; the destination half and the virial accumulate in registers.
+      // Lane c of a group adds component c, so the source atomics of an edge
+      // share one instruction and one sector once the group spans the three
+      // components.
+      const float* vector = args.edge_vec + edge * 3;
+#pragma unroll
+      for (int first = 0; first < 3; first += EdgeWidth) {
+        const int component = first + base_channel;
+        if (component < 3) {
+          const float value = vector_component(output, component);
+          unsigned long long* const source_words =
+              args.force + geometry.source * kForceWords;
+          check_force_range(value, source_words);
+          atomicAdd(source_words + component, fixed_point(-value));
+          destination_force[first / EdgeWidth] += value;
+        }
+      }
+#pragma unroll
+      for (int component = 0; component < 9; ++component) {
+        if (component % EdgeWidth == base_channel) {
+          node_virial[component / EdgeWidth] =
+              fmaf(-output[component / 3], vector[component % 3],
+                   node_virial[component / EdgeWidth]);
+        }
+      }
+    }
     if constexpr (HasSpin) {
       // The magnetic cotangent belongs to the source node, which this
-      // destination-major scan does not own. It is emitted per edge and
-      // reduced onto sources by the shared edge assembly, which already walks
-      // the source CSR for the conservative force.
+      // destination-major scan does not own. It is either emitted per edge
+      // and reduced onto sources by the shared edge assembly, or added onto
+      // its source by the fused assembly. Every contribution above
+      // differentiates the conditioned moment, so the source scale is the
+      // single remaining chain factor to the raw input moment.
 #pragma unroll
       for (int component = 0; component < 3; ++component) {
         spin_cotangent[component] =
-            subwarp_sum(spin_cotangent[component], mask);
+            neighbor_scale *
+            subwarp_allsum(spin_cotangent[component], mask);
       }
-      if (thread == leader) {
-        // Every contribution above differentiates the conditioned moment, so
-        // the store is the single point that applies the remaining chain
-        // factor and hands back a gradient of the raw input moment.
-        args.edge_spin_gradient[edge * 3 + 0] =
-            neighbor_scale * spin_cotangent[0];
-        args.edge_spin_gradient[edge * 3 + 1] =
-            neighbor_scale * spin_cotangent[1];
-        args.edge_spin_gradient[edge * 3 + 2] =
-            neighbor_scale * spin_cotangent[2];
+      if (thread == leader && args.edge_spin_gradient != nullptr) {
+        args.edge_spin_gradient[edge * 3 + 0] = spin_cotangent[0];
+        args.edge_spin_gradient[edge * 3 + 1] = spin_cotangent[1];
+        args.edge_spin_gradient[edge * 3 + 2] = spin_cotangent[2];
+      }
+      if (assemble) {
+#pragma unroll
+        for (int first = 0; first < 3; first += EdgeWidth) {
+          const int component = first + base_channel;
+          if (component < 3) {
+            const float value = vector_component(spin_cotangent, component);
+            check_force_range(value,
+                              args.force + geometry.source * kForceWords);
+            atomicAdd(
+                args.spin_force + geometry.source * kForceWords + component,
+                fixed_point(value));
+          }
+        }
+      }
+    }
+  };
+  // The first warp of a shared row requests the lines of the row for every
+  // warp of the row.
+  if constexpr (Canonical) {
+    if (part == 0) {
+      prefetch_row_bonds(args, begin, end, thread);
+      if constexpr (SlotScan) {
+        prefetch_row_sources(args, edge_index, begin, end, thread);
+      }
+    }
+  }
+  const auto pairwise_edge = [&](const long position) {
+    scan_edge(position, ScanTag{});
+  };
+  if constexpr (SlotScan) {
+    walk_row(
+        begin, end, group, thread, part, args.edge_vec, args.table_max,
+        [&](const long position) { scan_edge(position, ScanTag{}); });
+  } else if constexpr (Foldable) {
+    if (folded) {
+      walk_row(
+          begin, end, group, thread, part, args.edge_vec, args.table_max,
+          [&](const long position) { scan_edge(position, ScanTag{}); });
+    } else {
+      walk_row(begin, end, group, thread, part,
+                                         args.edge_vec, args.table_max,
+                                         pairwise_edge);
+    }
+  } else {
+    walk_row(begin, end, group, thread, part,
+                                       args.edge_vec, args.table_max,
+                                       pairwise_edge);
+  }
+  if (!active) {
+    return;
+  }
+  // The node totals of the warps that share the row, which `accumulators`
+  // visits in its order.
+  constexpr int Accumulators = ForceSlots + VirialSlots + (Bridged ? 1 : 0);
+  __shared__ float split_buffer[split_buffer_floats()];
+  const auto accumulators = [&](auto&& visit) {
+#pragma unroll
+    for (int slot = 0; slot < ForceSlots; ++slot) {
+      visit(destination_force[slot]);
+    }
+#pragma unroll
+    for (int slot = 0; slot < VirialSlots; ++slot) {
+      visit(node_virial[slot]);
+    }
+    if constexpr (Bridged) {
+      visit(pair_energy);
+    }
+  };
+  if (!merge_split_row(part, thread, split_buffer,
+                                            accumulators)) {
+    return;
+  }
+  if (assemble) {
+    // Every group holds partial totals of its own edges in the same slots,
+    // so one merge over the groups completes the node.
+#pragma unroll
+    for (int slot = 0; slot < ForceSlots; ++slot) {
+      destination_force[slot] =
+          reduce_channel_groups(destination_force[slot]);
+    }
+#pragma unroll
+    for (int slot = 0; slot < VirialSlots; ++slot) {
+      node_virial[slot] = reduce_channel_groups(node_virial[slot]);
+    }
+    if (thread < EdgeWidth) {
+      const long absolute = args.node_begin + node;
+#pragma unroll
+      for (int first = 0; first < 3; first += EdgeWidth) {
+        const int component = first + base_channel;
+        if (component < 3) {
+          const float value = destination_force[first / EdgeWidth];
+          check_force_range(value, args.force + absolute * kForceWords);
+          atomicAdd(args.force + absolute * kForceWords + component,
+                    fixed_point(value));
+        }
+      }
+#pragma unroll
+      for (int component = 0; component < 9; ++component) {
+        if (component % EdgeWidth == base_channel) {
+          args.node_virial[node * 9 + component] =
+              node_virial[component / EdgeWidth];
+        }
       }
     }
   }
@@ -1990,6 +4185,118 @@ __global__ void zero_padding_kernel(long node_count,
 
 // === Launch dispatch ===
 
+/// Launch the edge scans `Family::kernel` of blocks of `Rows`
+/// independent rows, or, for a canonical system of 32-bit indices whose rows
+/// `row_split` shares among `Split > 1` warps, the scans
+/// `Family::kernel<1, Split, Fold>`.
+template 
+void launch_row_scans(const Arguments& args, cudaStream_t stream) {
+  if constexpr (Canonical && std::is_same_v) {
+    switch (row_split, Rows>(
+        args.system_node_count)) {
+      case 8:
+        launch_edge_scan<1, 8>(Family::template kernel<1, 8, Fold>, args,
+                               stream);
+        return;
+      case 4:
+        launch_edge_scan<1, 4>(Family::template kernel<1, 4, Fold>, args,
+                               stream);
+        return;
+      case 2:
+        launch_edge_scan<1, 2>(Family::template kernel<1, 2, Fold>, args,
+                               stream);
+        return;
+      default:
+        break;
+    }
+  }
+  launch_edge_scan(Family::template kernel, args,
+                            stream);
+}
+
+/// Launch the co-scheduled blocks `Family::kernel` when the
+/// system fills the device with them, and return whether it did; a profile
+/// that does not co-schedule (`Warps == 1`) never does.
+template 
+bool launch_coscheduled_scans(const Arguments& args, cudaStream_t stream) {
+  if constexpr (Warps > 1) {
+    if (fills_coscheduled(args.system_node_count)) {
+      launch_edge_scan(Family::template kernel, args,
+                                 stream);
+      return true;
+    }
+  }
+  return false;
+}
+
+/// Launch the edge scans of one direction, whose kernels are the
+/// `Family::kernel` of the single-type scan `Fold` of the
+/// descriptor: co-scheduled blocks of `Warps` rows when the system fills the
+/// device with them, rows that a small canonical system shares among several
+/// warps, and blocks of independent rows otherwise, one for a spin-free
+/// profile and `kSlotRows` for a spin profile. Every launch also launches the
+/// slotted kernels of the profile (`Family::launch_slots`): the number
+/// of node types of the system, which the device knows, selects the kernel
+/// that scans the rows, and the blocks of the others retire.
+template 
+void launch_edge_scans(const Arguments& args, cudaStream_t stream) {
+  dispatch_fold(args, [&](auto fold) {
+    constexpr Scan Fold = decltype(fold)::value;
+    constexpr int Warps =
+        coscheduled_warps();
+    if (!launch_coscheduled_scans(args, stream)) {
+      launch_row_scans(args, stream);
+    }
+    Family::template launch_slots(args, stream);
+  });
+}
+
+/// The slotted scan of `Slots` slots of a descriptor whose single-type scan
+/// is `Fold`: the one that slots the basis functions of its low-rank table,
+/// or the one that reads the radial table.
+template 
+constexpr Scan slotted_scan() {
+  if constexpr (Slots > 4) {
+    return Fold == Scan::LowRank ? Scan::SlottedLowRank8 : Scan::Slotted8;
+  } else {
+    return Fold == Scan::LowRank ? Scan::SlottedLowRank4 : Scan::Slotted4;
+  }
+}
+
+/// Launch the slotted kernels `Family::kernel` of a
+/// launch that carries the bitmap of the node types, the launches whose
+/// systems they may serve: co-scheduled blocks of `slot_coscheduled_warps`
+/// rows when the system fills the device with them, rows shared among the
+/// warps of a block when the system leaves most of the resident warps of the
+/// slotted kernel idle (`row_split`), and blocks of `kSlotRows` independent
+/// rows otherwise.
+template 
+void launch_slot_scans(const Arguments& args, cudaStream_t stream) {
+  if (args.type_present == nullptr) {
+    return;
+  }
+  constexpr int Warps = slot_coscheduled_warps();
+  if (!launch_coscheduled_scans(args, stream)) {
+    launch_row_scans(args,
+                                                                     stream);
+  }
+}
+
 template 
 struct ForwardLauncher {
+  template 
+  static constexpr auto kernel = forward_kernel;
+
+  /// Launch the slotted forward scans of a descriptor whose single-type
+  /// scan is `Fold` (see `runs_slotted`).
+  template 
+  static void launch_slots(const Arguments& args, cudaStream_t stream) {
+    if constexpr (runs_slotted()) {
+      launch_slot_scans()>(args,
+                                                                     stream);
+    }
+    if constexpr (runs_slotted()) {
+      launch_slot_scans()>(args,
+                                                                     stream);
+    }
+  }
+
   static void run(const Arguments& args, cudaStream_t stream) {
     using P = Profile;
-    forward_kernel
+    launch_edge_scans(args, stream);
+    forward_readout_kernel
         <<(args.node_count), P::Threads, 0, stream>>>(args);
   }
 };
@@ -2027,31 +4365,67 @@ struct BackwardLauncher {
   // the device's warp capacity the wide variant wins by 2.1 % of the whole
   // backward, at a third the two are within noise, and above that the narrow
   // variant wins by up to 3.6 %. The threshold is therefore stated against the
-  // device's own warp capacity, and the query is made once per configuration.
+  // device's own warp capacity, and the query is made once per configuration
+  // and device.
   static bool prefer_wide_nodes() {
-    static const bool wide = [] {
+    static DeviceMemo preferences;
+    return preferences.get([](int device) {
       using Narrow = Profile;
       using Wide = Profile;
       int narrow_blocks = 0;
       int wide_blocks = 0;
-      cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &narrow_blocks,
-          node_backward_kernel,
-          Narrow::Threads, 0);
-      cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &wide_blocks,
-          node_backward_kernel,
-          Wide::Threads, 0);
-      int device = 0;
-      cudaGetDevice(&device);
+      check_query(
+          cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+              &narrow_blocks,
+              node_backward_kernel,
+              Narrow::Threads, 0),
+          "the occupancy query of the node backward");
+      check_query(
+          cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+              &wide_blocks,
+              node_backward_kernel,
+              Wide::Threads, 0),
+          "the occupancy query of the node backward");
       int threads_per_sm = 0;
-      cudaDeviceGetAttribute(&threads_per_sm,
-                             cudaDevAttrMaxThreadsPerMultiProcessor, device);
+      check_query(
+          cudaDeviceGetAttribute(
+              &threads_per_sm, cudaDevAttrMaxThreadsPerMultiProcessor, device),
+          "the thread capacity of a multiprocessor");
       const int warps_per_sm = threads_per_sm / kWarpSize;
       const int narrow_warps = narrow_blocks * (Narrow::Threads / kWarpSize);
-      return wide_blocks > narrow_blocks && narrow_warps * 4 < warps_per_sm;
-    }();
-    return wide;
+      return static_cast(wide_blocks > narrow_blocks &&
+                               narrow_warps * 4 < warps_per_sm);
+    }) != 0;
+  }
+
+  template 
+  static constexpr auto kernel = edge_backward_kernel;
+
+  /// Launch the slotted edge backwards of a descriptor whose single-type
+  /// scan is `Fold` (see `runs_slotted`).
+  template 
+  static void launch_slots(const Arguments& args, cudaStream_t stream) {
+    if constexpr (runs_slotted()) {
+      launch_slot_scans()>(args,
+                                                                     stream);
+    }
+    if constexpr (runs_slotted()) {
+      launch_slot_scans()>(args,
+                                                                     stream);
+    }
   }
 
   template 
@@ -2065,19 +4439,17 @@ struct BackwardLauncher {
   }
 
   static void run(const Arguments& args, cudaStream_t stream) {
-    using P = Profile;
     if (prefer_wide_nodes()) {
       launch_node_backward(args, stream);
     } else {
       launch_node_backward(args, stream);
     }
-    edge_backward_kernel
-        <<(args.node_count), P::Threads, 0, stream>>>(args);
+    launch_edge_scans(args, stream);
     // The reserved edge slots beyond the physical count are only known on the
     // device, so the grid is sized from the storage bound and the surplus
     // blocks retire immediately.
-    if (args.clear_padding) {
+    if (args.clear_padding && args.edge_gradient != nullptr) {
       constexpr int kPaddingThreads = 128;
       constexpr long kPaddingBlockLimit = 1024;
       const long padding_blocks = min(
@@ -2172,40 +4544,14 @@ void dispatch_modes(const Arguments& args, cudaStream_t stream) {
   }
 }
 
-// The angular degree is a compile-time specialization because degrees three
-// and four add moment accumulators that must not enter the register budget of
-// the production ``lmax=2`` path.
-//
-// The operator entry point validates the degree before it reaches this
-// dispatch. The unreachable default is still checked rather than folded into
-// the highest degree, so that a degree outside the compiled set can only ever
-// fail loudly instead of running a kernel for a different model.
-template  class L>
-void dispatch_degree(const Arguments& args, cudaStream_t stream) {
-  switch (args.lmax) {
-    case 2:
-      dispatch_modes(args, stream);
-      return;
-    case 3:
-      dispatch_modes(args, stream);
-      return;
-    case 4:
-      dispatch_modes(args, stream);
-      return;
-    default:
-      throw std::runtime_error("dpa4c: uncompiled angular degree");
-  }
+template 
+void launch_forward(const Arguments& arguments, cudaStream_t stream) {
+  dispatch_modes(arguments, stream);
 }
 
-#define DPA4C_DEFINE_CHANNEL(width)                              \
-  void launch_forward_c##width(const Arguments& arguments,       \
-                               cudaStream_t stream) {            \
-    dispatch_degree(arguments, stream);  \
-  }                                                              \
-  void launch_backward_c##width(const Arguments& arguments,      \
-                                cudaStream_t stream) {           \
-    dispatch_degree(arguments, stream); \
-  }
+template 
+void launch_backward(const Arguments& arguments, cudaStream_t stream) {
+  dispatch_modes(arguments, stream);
+}
 
 }  // namespace deepmd_dpa4c
diff --git a/source/op/pt/dpa4c/graph_compress_launch.h b/source/op/pt/dpa4c/graph_compress_launch.h
index 8c220f37d6..4ce4be6f30 100644
--- a/source/op/pt/dpa4c/graph_compress_launch.h
+++ b/source/op/pt/dpa4c/graph_compress_launch.h
@@ -2,10 +2,11 @@
 //
 // Host-side launch interface of the compressed DPA4C descriptor.
 //
-// The kernel templates are instantiated once per scalar width in a dedicated
-// translation unit, which keeps the compile time of the angular-degree and
-// topology specializations bounded and parallel. The dispatching translation
-// unit only sees the plain argument bundle declared here.
+// The kernel templates are instantiated in generated translation units
+// (shard.cu.in), one per angular degree for the widest scalar widths and one
+// per width for the others, which keeps the compile time of the topology
+// specializations bounded and parallel. The dispatching translation unit only
+// sees the plain argument bundle declared here.
 
 #pragma once
 
@@ -22,12 +23,29 @@ constexpr int kWarpSize = 32;
 /// non-negative range that node and edge identifiers occupy.
 enum class IndexKind : int { Bits32 = 0, Bits64 = 1 };
 
+/// Words of the presence bitmap of the node types of a system, which covers
+/// the types below 32 times this count (`Arguments::type_present`).
+constexpr int kTypeWords = 4;
+
+/// Mode profiles that a folded scan carries as extra features and that a
+/// slotted scan holds in the slots of its types. A model with more modes
+/// scans unfolded.
+constexpr int kFoldedModes = 4;
+
 /// Complete argument bundle of one descriptor launch.
 ///
 /// Pointers that a given direction does not read stay null. All tensors are
-/// contiguous and, except for the topology and the coupling layout, fp32.
+/// contiguous. The tables, weights and per-node and per-edge features are
+/// fp32; the topology and the coupling layout are integers, the node-type
+/// words an integer flag, an fp32 radius and a bitmap, the force totals
+/// 64-bit fixed point, and the pair energy fp64.
 struct Arguments {
   long node_count = 0;
+  // Node count of the whole system, from which the edge scans choose their
+  // row split and their co-scheduled blocks, both of which change the
+  // summation order of a row, so that every run of a tiled system takes the
+  // schedule of the whole system.
+  long system_node_count = 0;
   long edge_count = 0;
   int lmax = 2;
   int interval_count = 0;
@@ -43,6 +61,15 @@ struct Arguments {
   // last node may clear them; an earlier run would erase the gradients that
   // later runs are about to write.
   bool clear_padding = true;
+  // Whether the canonical edge scans request the bond vectors of each row into
+  // L2 as the row starts, and the slotted scans its source indices as well.
+  // Set when the bond vectors of the edge storage exceed the L2 cache, as
+  // those of a graph that fills the storage then stream from DRAM: a row's
+  // vectors are contiguous, the lanes of a warp that scans the row request
+  // one 128-byte line each until the row is covered, and the transfer then
+  // overlaps the scan of the row instead of stalling the first use of every
+  // line. A caller may size the storage beyond the edges.
+  bool prefetch_bonds = false;
   int coupling_count = 0;
   float table_stride = 0.0f;
   float table_max = 0.0f;
@@ -73,6 +100,15 @@ struct Arguments {
   const long* destination_row_ptr = nullptr;
   const long* atype = nullptr;
   const float* table = nullptr;
+  // Low-rank radial table of the single-type and slotted scans, null when
+  // the descriptor carries none: the leading ``C1`` channels, the ``Basis``
+  // functions on which the scalar-only channels are expanded, then the radial
+  // modes, in the layout of ``table``. ``basis_matrix`` is the transpose of
+  // the expansion, of shape ``(Basis, C0 - C1)``: channel ``c >= C1`` is the
+  // sum over ``k`` of ``basis_matrix[k * (C0 - C1) + c - C1]`` times basis
+  // function ``k``.
+  const float* basis_table = nullptr;
+  const float* basis_matrix = nullptr;
   const float* pair_film = nullptr;
   const float* pair_mixing = nullptr;
   const float* type_embedding = nullptr;
@@ -89,6 +125,17 @@ struct Arguments {
   const float* contact_radius = nullptr;
   const float* descriptor_gradient = nullptr;
   const float* state = nullptr;
+  // Device word that stays zero while every node of the system carries the
+  // type of node zero. A canonical launch of such a system scans the bare
+  // radial features; null disables that scan.
+  const int* type_mismatch = nullptr;
+  // Largest length-scale radius among the node types of the system, written
+  // on the device by the scan of the node types that precedes the edge kernels
+  // of a canonical launch, or null. No bond of a row lies inside a bridging
+  // window once it is at least `f_outer` times the sum of the centre radius
+  // and this radius long, so the edge scans resolve the window of a pair only
+  // for the shorter bonds of the row.
+  const float* largest_radius = nullptr;
 
   // === Native spin ===
   // Present together or not at all. ``spin`` is indexed by absolute node id,
@@ -116,14 +163,81 @@ struct Arguments {
   float* descriptor = nullptr;
   float* state_out = nullptr;
   float* moment_gradient = nullptr;
+  // Per-edge energy cotangent. A null pointer skips the store, which a
+  // backward that assembles the force itself only needs for the per-atom
+  // virial.
   float* edge_gradient = nullptr;
   // Per-node on-site magnetic gradient and the per-edge contribution to the
   // source node's magnetic gradient. The latter is reduced onto source nodes
-  // by the shared edge force assembly, which already walks the source CSR.
+  // either by the shared edge force assembly, which walks the source CSR, or
+  // by the fused assembly below.
   float* spin_gradient = nullptr;
   float* edge_spin_gradient = nullptr;
+
+  // === Fused force assembly ===
+  // Present together or not at all. The edge backward then closes the
+  // conservative force and the virial itself:
+  //
+  //   force[i] = sum(dst=i) g_e - sum(src=i) g_e
+  //   W_i      = sum(dst=i) -g_e (x) r_e
+  //
+  // ``force`` and, for a spin-conditioned launch, ``spin_force`` (the source
+  // totals of the magnetic edge cotangent) hold signed fixed-point sums in
+  // units of ``1 / kForceFixedScale`` over the absolute node axis, like
+  // ``atype``, because the source half lands on nodes of other runs; node
+  // ``n`` owns the ``kForceWords`` words from ``n * kForceWords``.
+  // ``node_virial`` holds the nine destination-attributed components of each
+  // node of the run, offset by the caller like ``descriptor``.
+  unsigned long long* force = nullptr;
+  unsigned long long* spin_force = nullptr;
+  float* node_virial = nullptr;
+
+  // === Node types of the system ===
+  // Bitmap of the node types of the system, `kTypeWords` device words in
+  // which bit `t % 32` of word `t / 32` is set when a node carries type `t`,
+  // written by the scan of the node types that also writes `type_mismatch`.
+  // A canonical launch of a system of a few types slots the mode profiles
+  // and the low-rank basis functions of every edge by the type of its source
+  // (see `TypeSlots`). It is null, which disables the slotted scans, for a
+  // launch that is not canonical, of more types than the bitmap covers, the
+  // padding type included, of more than `kFoldedModes` radial modes, and of a
+  // descriptor without modes and without a low-rank table.
+  const unsigned* type_present = nullptr;
 };
 
+/// Resolution of the fixed-point force sums, 2^32 per unit.
+///
+/// Integer addition is associative, so the totals do not depend on the order
+/// in which concurrent atomics land and the force is reproducible bit for bit.
+/// Each contribution is rounded to the nearest multiple of 2^-32 eV/Å, an
+/// absolute error of at most 1.2e-10 eV/Å that lies below its fp32 rounding
+/// once it exceeds about 2e-3 eV/Å, and a node keeps a range of 2^31 eV/Å
+/// (see kForceFixedLimit).
+constexpr float kForceFixedScale = 4294967296.0f;
+
+/// Fixed-point words per node of the force sums: the three components and the
+/// range word, so that the 32 bytes of a node fill exactly one sector. The
+/// three atomics an edge adds onto its source then land in one sector and
+/// travel in one instruction, which is what bounds the source scatter of the
+/// edge backward.
+constexpr int kForceWords = 4;
+
+/// Word of the fixed-point totals of a node that a contribution of
+/// kForceFixedLimit or more sets; the node then reports its energy as not a
+/// number, which marks the forces of its frame.
+constexpr int kForceRangeWord = 3;
+
+/// Magnitude in eV/Å from which a force contribution lies outside the range
+/// of the fixed-point totals: half the range of a node, which every
+/// contribution whose conversion would saturate reaches, while any two
+/// contributions below it sum within the range. The check covers each
+/// contribution: a total that leaves the range through three or more
+/// contributions below the limit is not detected, which takes three
+/// neighbours of one atom within 0.001 Å (carbon) to 0.01 Å (uranium) of it
+/// under the unscreened Coulomb repulsion, and closer under any screened pair
+/// term.
+constexpr float kForceFixedLimit = 1073741824.0f;
+
 // === Compile-time descriptor profile ===
 
 constexpr int triangular(int width) { return width * (width + 1) / 2; }
@@ -142,6 +256,28 @@ constexpr int degree_two_width(int channels) {
   return width < 4 ? 4 : width;
 }
 
+// Functions on which the single-type and slotted scans expand the
+// scalar-only channels, the channels from ``C1`` on that feed degree zero
+// alone.
+//
+// The radial map of a trained descriptor spans far fewer independent functions
+// than it has channels: on the released checkpoints the largest singular value
+// of its scalar-only channels that a 32-function basis discards lies more than
+// nine orders of magnitude below the leading one, and the basis reproduces
+// every channel and its radial derivative to the rounding of the
+// single-precision table. A scan that accumulates the basis functions instead
+// of the channels reads that many table columns per edge instead of
+// ``C0 - C1``, and the expansion runs once per node, or once per node type of
+// the system in a slotted scan. The width is a multiple of every edge-group
+// width of these profiles, so the basis columns tile the lanes evenly.
+// Narrower profiles have too few scalar-only channels for an expansion
+// to save reads. The compression adopts the basis only when it reproduces the
+// channels of the descriptor at hand, so a model of higher rank keeps the
+// direct scan.
+constexpr int low_rank_basis(int channels) {
+  return channels == 64 || channels == 128 ? 32 : 0;
+}
+
 constexpr int degree_rank(int degree, int rank_one, int rank_two) {
   return degree == 1 ? rank_one : (degree == 2 ? rank_two : 1);
 }
@@ -284,6 +420,36 @@ struct EdgeMap<128> {
   static constexpr int SpinBackward = 8;
 };
 
+// Co-scheduled rows of the canonical edge scans.
+//
+// An edge reads one row of the radial table, and the row is served from L1
+// only while the warps resident on a multiprocessor read nearby rows at the
+// same time. A destination row is sorted by bond length, so a warp sweeps the
+// table from short to long bonds, and warps that start their rows together
+// stay in step. One warp per block lets the block scheduler stagger the
+// resident warps until they read unrelated rows, and a wide table then
+// streams from L2 once per edge. The scans of the wide profiles therefore
+// give one block `Warps` consecutive destination rows, enough warps to fill a
+// multiprocessor on its own, and realign them at bond-length intervals of
+// `kLockstepWindow` with a block barrier. The narrow profiles hold their
+// table in L1 under any schedule and keep blocks of independent rows, as does
+// every system too small to give each multiprocessor two co-scheduled blocks,
+// unless the split scans share its rows among several warps (`row_split`).
+// The widths are the measured optima on disordered diamond carbon; the
+// backward of the high-degree profiles needs the realignment one profile
+// earlier because its edges carry more work between table reads. A spin
+// profile realigns where its spin-free counterpart does, in blocks of fewer
+// rows (`coscheduled_warps`).
+template 
+struct Lockstep {
+  static constexpr int Warps =
+      Channels >= 128 || (Backward && Channels >= 64 && Lmax >= 3) ? 24 : 1;
+};
+
+/// Bond-length interval in Å between two realignments of the co-scheduled
+/// rows of one block.
+constexpr float kLockstepWindow = 4.0f;
+
 // ``HasSpin`` is deliberately without a default. It changes the moment
 // layout and the descriptor width, so an instantiation that omits it would
 // silently read a spin-free layout out of a spin-conditioned buffer.
@@ -297,6 +463,10 @@ struct Profile {
   static constexpr int C2 = degree_two_width(Channels);
   static constexpr int K1 = C2;
   static constexpr int K2 = 2;
+  // Width of the low-rank expansion of the scalar-only channels, zero for a
+  // profile whose single-type scan reads the channels directly. The spin
+  // branch never folds, so a spin-conditioned profile carries none.
+  static constexpr int Basis = HasSpin ? 0 : low_rank_basis(Channels);
 
   // === Native spin widths ===
   // The neighbour spin width is derived from the degree-two width, so the
@@ -394,27 +564,16 @@ struct Profile {
   static constexpr int Threads = kWarpSize;
 };
 
-/// Resident blocks an edge kernel is compiled for, passed as the minimum block
-/// count of its launch bounds.
-///
-/// A block is one warp, so the count sets the register budget per thread:
-/// thirty-two blocks exhaust the 65,536-register file at sixty-four registers
-/// per thread, and a lower count raises the budget in proportion at the cost of
-/// occupancy. Zero sets no budget and leaves the register count to the
-/// compiler.
-template 
+/// Resident warps per multiprocessor an edge kernel is compiled for. A kernel
+/// of blocks of W warps passes the count divided by W as the minimum block
+/// count of its launch bounds, so the count sets the register budget per
+/// thread: thirty-two warps exhaust the 65,536-register file at sixty-four
+/// registers per thread, and a lower count raises the budget in proportion at
+/// the cost of occupancy. Zero sets no budget and leaves the register count to
+/// the compiler. The comments below count one-warp blocks, which a
+/// multiprocessor holds as many of as its block limit allows.
+template 
 constexpr int edge_resident_blocks() {
-  // Measured on compute capability 9.0, degree two holds its geometric working
-  // set in sixty-four registers and spills only what the five spin families
-  // add, so twenty-four blocks buy eighty registers, remove the spill outright
-  // and are worth 2.0% to 4.5% of the step. Degree three overflows the budget
-  // on geometry alone -- it spills without the spin branch and still spills at
-  // eighty and at ninety-six registers -- so the occupancy it would give up
-  // buys an incomplete fix and costs 7.4% to 9.6%. The relief is therefore
-  // taken only where it is complete.
-  if (HasSpin && Lmax == 2) {
-    return 24;
-  }
 #if defined(__CUDA_ARCH__) && \
     (__CUDA_ARCH__ == 890 || (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300))
   // A multiprocessor of compute capability 8.9 or 12.x holds at most
@@ -426,24 +585,112 @@ constexpr int edge_resident_blocks() {
   // compiler's own count, which keeps the spin families in registers and suits
   // the narrow profiles, whose two-lane edge groups keep sixteen edges in
   // flight per warp.
+  //
+  // The spin kernels keep the budget where their working set fits eighty
+  // registers, which on 12.0 holds for the degree-two forward kernels and for
+  // the degree-two backward kernels of sixty-four channels and more. The
+  // others, whose geometric state at degree three or whose spin payload on
+  // lanes that own few channels exceeds eighty registers, run fastest on the
+  // compiler's own count.
+  if (HasSpin) {
+    return Lmax == 2 && (!Backward || Channels >= 64) ? 24 : 0;
+  }
   return HasModes && Channels >= 64 ? 24 : 0;
+#elif defined(__CUDA_ARCH__) && \
+    (__CUDA_ARCH__ == 750 || __CUDA_ARCH__ == 860 || __CUDA_ARCH__ == 870)
+  // A multiprocessor of compute capability 7.5, 8.6 or 8.7 holds at most
+  // sixteen blocks, which leaves a one-warp block 128 registers, more than
+  // any edge kernel takes, so the one-warp kernels keep the compiler's own
+  // count there. A co-scheduled block of twenty-four warps is bounded by its
+  // thread count alone, at eighty registers, and one of twenty, at ninety-six,
+  // on every architecture.
+  return 0;
+#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 700
+  // At the thirty-two resident blocks that a multiprocessor of compute
+  // capability 7.0 holds, sixty-four registers per thread, the edge backward of
+  // every spin-free profile spills 40 to 224 bytes. Twenty-four blocks, eighty
+  // registers, remove all but 8 to 24 bytes of that spill at degree two and
+  // about halve it at degree three. On a V100 they shorten the LAMMPS step of
+  // Nano by 13% to 30% on carbon, BCN and five elements alike, and that of Mini
+  // by 2% to 9% on BCN and five elements and by up to 35% on carbon, whose
+  // single-type scan keeps the sums of its extra feature live through the edge
+  // loop; against the compiler's own count they shorten the Nano step by 1.5%
+  // to 5.4% on BCN and five elements.
+  //
+  // The spin kernels of degree two take twenty-four blocks for the forward and
+  // for the backward of sixty-four channels and more, and the compiler's own
+  // count for the backward of fewer channels, at which the spin backward of
+  // Nano and Mini holds its payload without spilling and the spin step of Nano
+  // is 14% to 22% shorter and that of Mini up to 1%. The kernels of degree
+  // three serve the generic graph path and the LAMMPS systems of more than
+  // eight node types; unmeasured on 7.0, they keep the thirty-two blocks of the
+  // architectures without a branch of their own.
+  if (HasSpin) {
+    return Lmax == 2 ? (!Backward || Channels >= 64 ? 24 : 0) : 32;
+  }
+  return 24;
 #else
-  return 32;
+  // Measured on compute capability 9.0, whose multiprocessor holds thirty-two
+  // blocks; 8.0 and 10.x take these values unmeasured. Thirty-two blocks leave
+  // sixty-four registers per thread. The backward of eight channels, whose
+  // two-lane edge groups keep sixteen edges in flight per warp, spills 56
+  // bytes there; twenty-four blocks give it eighty registers, which remove the
+  // spill, shorten the kernel by 14% and the LAMMPS step of such a model by
+  // 3.3% to 5.3% at 10,296 and 100,800 atoms. The backward of sixteen channels
+  // does not spill at sixty-four registers and takes 3% to 5% longer at
+  // eighty, those of thirty-two and sixty-four channels take up to 5% longer
+  // at eighty on a system of one element, and the forward kernels change by at
+  // most 4% either way, under 1% of the step, so the other kernels keep
+  // thirty-two blocks. The degree-two spin backward spills 16 to 104 bytes at
+  // the eighty registers of twenty-four blocks and runs 3% to 6% faster on the
+  // compiler's own count, 96 to 127 registers, while its forward runs fastest
+  // at twenty-four blocks. The degree-three spin kernels keep thirty-two
+  // blocks; a system of at most eight types scans them slotted, so no
+  // production path has measured them.
+  if (HasSpin) {
+    return Lmax == 2 ? (Backward ? 0 : 24) : 32;
+  }
+  return Backward && Channels == 8 ? 24 : 32;
 #endif
 }
 
-/// Scalar widths that own a compiled specialization.
+/// Resident warps per multiprocessor a slotted edge kernel is compiled for
+/// (see `edge_resident_blocks`), for blocks of `CoscheduledWarps` co-scheduled
+/// rows, or zero for blocks of independent or shared rows. The slots add one
+/// set of sums per slot to the state of a lane. A co-scheduled block fills a
+/// multiprocessor alone, so its rows are the resident warps, and its thread
+/// count sets the register budget. The other scans that slot the mode
+/// profiles keep sixteen resident warps, which the 65,536-register file of
+/// every supported architecture gives 128 registers per thread, and those
+/// that slot the basis functions alone keep twenty, 96 registers per thread;
+/// a block of one row shared by eight warps, a width that does not divide
+/// twenty, keeps sixteen there as well.
+template 
+constexpr int slot_resident_warps() {
+  if constexpr (CoscheduledWarps > 0) {
+    return CoscheduledWarps;
+  } else {
+    return HasModes ? 16 : 20;
+  }
+}
+
+/// Scalar widths that own a compiled specialization. Each width, with every
+/// degree that graph_compress.cu dispatches, owns instantiations in a shard
+/// that source/op/pt/CMakeLists.txt generates; a width or degree is added in
+/// both places, and a Linux build that misses one fails at the link of the
+/// operator library.
 #define DPA4C_FOR_EACH_CHANNEL(macro) \
   macro(8) macro(16) macro(32) macro(64) macro(128)
 
-#define DPA4C_DECLARE_CHANNEL(width)                        \
-  void launch_forward_c##width(const Arguments& arguments,  \
-                               cudaStream_t stream);        \
-  void launch_backward_c##width(const Arguments& arguments, \
-                                cudaStream_t stream);
-
-DPA4C_FOR_EACH_CHANNEL(DPA4C_DECLARE_CHANNEL)
-
-#undef DPA4C_DECLARE_CHANNEL
+/// One direction of the descriptor for one scalar width and angular degree.
+/// The angular degree is a compile-time specialization because degrees three
+/// and four add moment accumulators that must not enter the register budget
+/// of the production ``lmax=2`` path. Every width of DPA4C_FOR_EACH_CHANNEL
+/// and every degree from 2 to 4 is instantiated, both directions of a degree
+/// in one generated translation unit (shard.cu.in).
+template 
+void launch_forward(const Arguments& arguments, cudaStream_t stream);
+template 
+void launch_backward(const Arguments& arguments, cudaStream_t stream);
 
 }  // namespace deepmd_dpa4c
diff --git a/source/op/pt/dpa4c/ops.cc b/source/op/pt/dpa4c/ops.cc
index 4569bb9bbe..c249513fbe 100644
--- a/source/op/pt/dpa4c/ops.cc
+++ b/source/op/pt/dpa4c/ops.cc
@@ -17,13 +17,37 @@
 // and source indices only.
 //
 // The two operators that close an energy-force evaluation, the generic
-// backward and the fused canonical energy-gradient, also evaluate an
-// analytical pair potential when `pair_table` is not empty. Its energy and
-// radial slope share their exponentials and need no saved state, so the edge
-// scan that emits the edge gradient produces both: it adds
+// backward and the fused canonical energy-force, also evaluate an analytical
+// pair potential when `pair_table` is not empty. Its energy and radial slope
+// share their exponentials and need no saved state, so the edge scan that
+// differentiates the descriptor produces both: it adds
 // `pair_seed_i * dV/dr / 2` to the radial cotangent of every edge of node `i`
-// and accumulates `V / 2` into the node energy (`pair_energy`, or the
-// `energy` of the fused operator, whose `seed` is the pair seed).
+// and accumulates `V / 2` into the node energy (`pair_energy`, or the node
+// energy of the fused operator, whose pair seed is its ownership mask).
+//
+// The fused canonical operator derives the ownership mask from `n_node` and
+// `n_local` -- a node contributes energy when its type is real and it is among
+// the leading `n_local` nodes of its frame -- and returns the frame energy, the
+// masked node energy, the force, the frame virial, the per-atom virial, the
+// magnetic force of a spin-conditioned descriptor (empty otherwise) and the
+// mask itself. It computes the per-atom virial when it receives the source
+// CSR; a two-entry placeholder in place of `source_row_ptr` and
+// `source_order` declines it, which leaves it zero.
+//
+// The fused long-range operator evaluates a model whose atomic energy adds a
+// latent-charge long-range module with a charge network of one channel. It
+// takes the arguments of the fused operator up to `tile`, the network that
+// predicts the charges from the descriptor (`charge_ws` to `charge_act`, in
+// the layout of the energy fitting, with the per-type bias `charge_bias`),
+// the inputs of deepmd::latent_charge_sum (`coord` to `spacing`, see
+// source/op/pt/long_range/ops.cc), whose charged nodes are the owned real
+// nodes, and the fitting routes. It returns the outputs of the fused operator
+// but the magnetic force: the descriptor carries no spin condition and the
+// model no pair potential.
+//
+// Every schema carries, next to the radial table, the low-rank table and the
+// expansion that a single-type scan reads instead of the scalar-only channels
+// (empty together when the descriptor admits none; see `low_rank_basis`).
 //
 // Every schema carries the zone-bridging window as the two dimensionless
 // fractions `f_inner` and `f_outer` of a pair's own length scale, together
@@ -37,7 +61,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) {
       "dpa4c_graph_compress(Tensor edge_vec, Tensor edge_index, "
       "Tensor edge_mask, Tensor destination_order, "
       "Tensor destination_row_ptr, Tensor atype, Tensor table, "
-      "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, "
+      "Tensor basis_table, Tensor basis_matrix, Tensor pair_film, "
+      "Tensor pair_mixing, Tensor type_embedding, "
       "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, "
       "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, "
       "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, "
@@ -48,7 +73,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) {
       "dpa4c_graph_compress_backward(Tensor descriptor_gradient, "
       "Tensor state, Tensor edge_vec, Tensor edge_index, Tensor edge_mask, "
       "Tensor destination_order, Tensor destination_row_ptr, Tensor atype, "
-      "Tensor table, Tensor pair_film, Tensor pair_mixing, "
+      "Tensor table, Tensor basis_table, Tensor basis_matrix, "
+      "Tensor pair_film, Tensor pair_mixing, "
       "Tensor type_embedding, Tensor readout_matrices, Tensor coupling_meta, "
       "Tensor coupling_entry, Tensor coupling_value, Tensor output_mean, "
       "Tensor output_inv_std, Tensor spin, Tensor spin_pair, "
@@ -61,7 +87,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) {
   library.def(
       "dpa4c_canonical_compress(Tensor edge_vec, Tensor source, "
       "Tensor destination_row_ptr, Tensor atype, Tensor table, "
-      "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, "
+      "Tensor basis_table, Tensor basis_matrix, Tensor pair_film, "
+      "Tensor pair_mixing, Tensor type_embedding, "
       "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, "
       "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, "
       "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, "
@@ -72,7 +99,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) {
       "dpa4c_canonical_compress_backward(Tensor descriptor_gradient, "
       "Tensor state, Tensor edge_vec, Tensor source, "
       "Tensor destination_row_ptr, Tensor atype, Tensor table, "
-      "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, "
+      "Tensor basis_table, Tensor basis_matrix, Tensor pair_film, "
+      "Tensor pair_mixing, Tensor type_embedding, "
       "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, "
       "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, "
       "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, "
@@ -84,7 +112,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) {
       "dpa4c_canonical_compress_backward_inplace("
       "Tensor descriptor_gradient, Tensor(a!) state, Tensor edge_vec, "
       "Tensor source, Tensor destination_row_ptr, Tensor atype, Tensor table, "
-      "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, "
+      "Tensor basis_table, Tensor basis_matrix, Tensor pair_film, "
+      "Tensor pair_mixing, Tensor type_embedding, "
       "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, "
       "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, "
       "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, "
@@ -93,8 +122,10 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) {
       "-> (Tensor edge_gradient, Tensor spin_gradient, "
       "Tensor edge_spin_gradient)");
   library.def(
-      "dpa4c_canonical_compress_energy_gradient(Tensor edge_vec, "
-      "Tensor source, Tensor destination_row_ptr, Tensor atype, Tensor table, "
+      "dpa4c_canonical_compress_energy_force(Tensor edge_vec, "
+      "Tensor source, Tensor destination_row_ptr, Tensor source_row_ptr, "
+      "Tensor source_order, Tensor n_node, Tensor n_local, Tensor atype, "
+      "Tensor table, Tensor basis_table, Tensor basis_matrix, "
       "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, "
       "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, "
       "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, "
@@ -103,7 +134,28 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) {
       "float degree_floor, float f_inner, float f_outer, "
       "Tensor[] ws, Tensor[] bs, int[] resnets, "
       "Tensor w_head, Tensor b_head, Tensor bias_atom_e, int act, "
-      "Tensor seed, int tile, Tensor pair_table) "
-      "-> (Tensor energy, Tensor edge_gradient, Tensor spin_gradient, "
-      "Tensor edge_spin_gradient)");
+      "int tile, Tensor pair_table, int[]? fitting_routes=None) "
+      "-> (Tensor energy, Tensor atom_energy, Tensor force, Tensor virial, "
+      "Tensor atom_virial, Tensor force_mag, Tensor mask)");
+  library.def(
+      "dpa4c_canonical_compress_long_range_energy_force(Tensor edge_vec, "
+      "Tensor source, Tensor destination_row_ptr, Tensor source_row_ptr, "
+      "Tensor source_order, Tensor n_node, Tensor n_local, Tensor atype, "
+      "Tensor table, Tensor basis_table, Tensor basis_matrix, "
+      "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, "
+      "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, "
+      "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, "
+      "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, "
+      "int lmax, float table_stride, float table_max, float rcut, float eps, "
+      "float degree_floor, float f_inner, float f_outer, "
+      "Tensor[] ws, Tensor[] bs, int[] resnets, "
+      "Tensor w_head, Tensor b_head, Tensor bias_atom_e, int act, int tile, "
+      "Tensor[] charge_ws, Tensor[] charge_bs, int[] charge_resnets, "
+      "Tensor charge_w_head, Tensor charge_b_head, Tensor charge_bias, "
+      "int charge_act, Tensor coord, Tensor total_charge, Tensor? box, "
+      "Tensor coefficient, Tensor communicator, int kernel, int channel, "
+      "float prefactor, int order, float spacing, "
+      "int[]? fitting_routes=None, int[]? charge_fitting_routes=None) "
+      "-> (Tensor energy, Tensor atom_energy, Tensor force, Tensor virial, "
+      "Tensor atom_virial, Tensor mask)");
 }
diff --git a/source/op/pt/dpa4c/shard.cu.in b/source/op/pt/dpa4c/shard.cu.in
new file mode 100644
index 0000000000..db5fe904ef
--- /dev/null
+++ b/source/op/pt/dpa4c/shard.cu.in
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// Generated instantiation shard of the compressed DPA4C descriptor. The
+// maintained kernels live in graph_compress_kernel.cuh; CMake selects the
+// scalar width and the angular degrees of each source so that nvcc compiles
+// the template grid in parallel (see the shard layout in CMakeLists.txt).
+
+#include "@DPA4C_KERNEL_HEADER@"
+
+namespace deepmd_dpa4c {
+
+@DPA4C_SHARD_INSTANTIATIONS@
+}  // namespace deepmd_dpa4c
diff --git a/source/op/pt/edge_force_virial.cu b/source/op/pt/edge_force_virial.cu
index 3a5454e064..31f630ee2e 100644
--- a/source/op/pt/edge_force_virial.cu
+++ b/source/op/pt/edge_force_virial.cu
@@ -11,9 +11,11 @@
 // Every node therefore writes its force and atom virial exactly once; the hot
 // path contains no global floating-point atomics. Per-frame virials are reduced
 // from the per-node values through two FP64 stages and cast to the model
-// precision only at the output boundary.
+// precision only at the output boundary. A caller that assembles the force
+// elsewhere reduces the source list alone for the atom virial.
 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -23,6 +25,8 @@
 #include 
 #include 
 
+#include "graph_ops.h"
+
 namespace {
 
 constexpr int kThreads = 256;
@@ -98,7 +102,7 @@ __global__ void edge_force_virial_kernel(
   scalar_t spin_z = 0;
   scalar_t virial[9] = {};
 
-  if (node < node_count) {
+  if (node < node_count && force != nullptr) {
     const long destination_begin = destination_row_ptr[node];
     const long destination_end = destination_row_ptr[node + 1];
     for (long position = destination_begin + lane; position < destination_end;
@@ -113,7 +117,8 @@ __global__ void edge_force_virial_kernel(
       destination_y += edge_gradient[edge * 3 + 1];
       destination_z += edge_gradient[edge * 3 + 2];
     }
-
+  }
+  if (node < node_count) {
     const long source_begin = source_row_ptr[node];
     const long source_end = source_row_ptr[node + 1];
     for (long position = source_begin + lane; position < source_end;
@@ -169,9 +174,11 @@ __global__ void edge_force_virial_kernel(
   }
 
   if (lane == 0 && node < node_count) {
-    force[node * 3 + 0] = destination_x - source_x;
-    force[node * 3 + 1] = destination_y - source_y;
-    force[node * 3 + 2] = destination_z - source_z;
+    if (force != nullptr) {
+      force[node * 3 + 0] = destination_x - source_x;
+      force[node * 3 + 1] = destination_y - source_y;
+      force[node * 3 + 2] = destination_z - source_z;
+    }
     scalar_t* output = node_virial + node * 9;
 #pragma unroll
     for (int component = 0; component < 9; ++component) {
@@ -192,6 +199,7 @@ __global__ void edge_force_virial_kernel(
 // fp64 whatever the stored type.
 template 
 __global__ void reduce_frame_segment_kernel(
+    long node_count,
     long frame_count,
     int partial_count,
     const long* __restrict__ frame_row_ptr,
@@ -205,8 +213,10 @@ __global__ void reduce_frame_segment_kernel(
     const long output = task / partial_count;
     const long frame = output / kComponents;
     const int component = output % kComponents;
-    const long begin = frame_row_ptr[frame];
-    const long end = frame_row_ptr[frame + 1];
+    // Frames are bounded by the node axis, so frame counts that sum past it
+    // never read beyond the values.
+    const long begin = min(frame_row_ptr[frame], node_count);
+    const long end = min(frame_row_ptr[frame + 1], node_count);
     double sum = 0.0;
     for (long node = begin + partial_index * static_cast(blockDim.x) +
                      threadIdx.x;
@@ -256,15 +266,6 @@ __global__ void finalize_frame_segment_kernel(
   }
 }
 
-// Exclusive prefix of the per-frame node counts, giving each frame's span of
-// the node axis as ``[row_ptr[f], row_ptr[f + 1])``.
-torch::Tensor frame_row_pointer(const torch::Tensor& n_node_per_frame) {
-  return torch::cat({torch::zeros({1}, n_node_per_frame.options()),
-                     torch::cumsum(n_node_per_frame, 0)})
-      .to(torch::kInt64)
-      .contiguous();
-}
-
 // Slices per frame that the partial pass reduces in parallel. Bounding it
 // keeps the partial buffer small for a single large frame, which is the
 // molecular dynamics case.
@@ -276,14 +277,14 @@ int frame_segment_partials(long node_count, long frame_count) {
 }
 
 // Sum ``node_values``, laid out as (node, component), over the node segment of
-// each frame. Returns a flat (frame * kComponents) buffer.
-template 
+// each frame into a flat (frame * kComponents) buffer of ``output_t``.
+template 
 void launch_frame_segment_sum(long node_count,
                               long frame_count,
                               const torch::Tensor& frame_row_ptr,
                               const scalar_t* node_values,
                               torch::Tensor& partial,
-                              scalar_t* out,
+                              output_t* out,
                               cudaStream_t stream) {
   const int partial_count = static_cast(partial.size(1));
   const long output_count = frame_count * kComponents;
@@ -291,12 +292,13 @@ void launch_frame_segment_sum(long node_count,
   const int partial_blocks = std::min(output_count * partial_count, 65535L);
   reduce_frame_segment_kernel
       <<>>(
-          frame_count, partial_count, frame_row_ptr.data_ptr(),
-          node_values, partial.data_ptr());
+          node_count, frame_count, partial_count,
+          frame_row_ptr.data_ptr(), node_values,
+          partial.data_ptr());
   FORCE_CHECK_LAUNCH("frame segment sum partial reduction");
 
   const int final_blocks = std::min(output_count, 65535L);
-  finalize_frame_segment_kernel
+  finalize_frame_segment_kernel
       <<>>(output_count, partial_count,
                                               partial.data_ptr(), out);
   FORCE_CHECK_LAUNCH("frame segment sum final reduction");
@@ -413,6 +415,13 @@ assemble_force_virial(long node_count,
 
 }  // namespace
 
+torch::Tensor frame_row_pointer(const torch::Tensor& n_node_per_frame) {
+  return torch::cat({torch::zeros({1}, n_node_per_frame.options()),
+                     torch::cumsum(n_node_per_frame, 0)})
+      .to(torch::kInt64)
+      .contiguous();
+}
+
 std::tuple
 build_graph_csr(torch::Tensor edge_index,
                 c10::SymInt node_count_symbol,
@@ -546,6 +555,11 @@ canonical_edge_force_virial(torch::Tensor edge_gradient,
                   source_row_ptr.numel() == node_count + 1,
               "canonical_edge_force_virial: row pointers must have N + 1 "
               "entries");
+  TORCH_CHECK(
+      source_order.dim() == 1 && source_order.numel() == edge_gradient.size(0),
+      "canonical_edge_force_virial: source_order must hold one "
+      "position per edge slot, expected ",
+      edge_gradient.size(0), ", got ", source_order.numel());
   TORCH_CHECK(
       valid_spin_cotangent(edge_gradient, edge_spin_gradient),
       "canonical_edge_force_virial: edge_spin_gradient must match the "
@@ -560,6 +574,79 @@ canonical_edge_force_virial(torch::Tensor edge_gradient,
                                edge_spin_gradient, want_atom_virial);
 }
 
+torch::Tensor canonical_source_virial(const torch::Tensor& edge_gradient,
+                                      const torch::Tensor& edge_vec,
+                                      const torch::Tensor& source_row_ptr,
+                                      const torch::Tensor& source_order,
+                                      long node_count) {
+  TORCH_CHECK(edge_gradient.scalar_type() == torch::kFloat32 &&
+                  edge_vec.scalar_type() == torch::kFloat32 &&
+                  source_order.scalar_type() == torch::kInt64 &&
+                  source_row_ptr.scalar_type() == torch::kInt64,
+              "canonical_source_virial: fp32 edges, int64 order and int64 "
+              "row pointers expected");
+  const auto resident = [&](const torch::Tensor& tensor) {
+    return tensor.is_cuda() &&
+           tensor.get_device() == c10::cuda::current_device() &&
+           tensor.is_contiguous();
+  };
+  TORCH_CHECK(resident(edge_gradient) && resident(edge_vec) &&
+                  resident(source_row_ptr) && resident(source_order),
+              "canonical_source_virial: every tensor must be contiguous and "
+              "on the current CUDA device");
+  TORCH_CHECK(edge_gradient.dim() == 2 && edge_gradient.size(1) == 3 &&
+                  edge_vec.sizes() == edge_gradient.sizes() &&
+                  source_order.numel() == edge_gradient.size(0) &&
+                  source_row_ptr.numel() == node_count + 1,
+              "canonical_source_virial: expected (S, 3) edges, S positions "
+              "and ",
+              node_count + 1, " row pointers, got ", edge_gradient.sizes(),
+              ", ", source_order.numel(), " and ", source_row_ptr.numel());
+  auto atom_virial = torch::empty({node_count, 3, 3}, edge_gradient.options());
+  if (node_count == 0) {
+    return atom_virial;
+  }
+  const int node_blocks =
+      static_cast((node_count + kWarpsPerBlock - 1) / kWarpsPerBlock);
+  edge_force_virial_kernel
+      <<>>(
+          node_count, edge_gradient.data_ptr(),
+          edge_vec.data_ptr(), nullptr, nullptr, nullptr,
+          source_order.data_ptr(), source_row_ptr.data_ptr(),
+          nullptr, nullptr, atom_virial.data_ptr(), nullptr);
+  FORCE_CHECK_LAUNCH("canonical_source_virial");
+  return atom_virial;
+}
+
+torch::Tensor frame_virial_sum(const torch::Tensor& node_virial,
+                               const torch::Tensor& n_node_per_frame) {
+  TORCH_CHECK(node_virial.scalar_type() == torch::kFloat32 &&
+                  node_virial.is_contiguous() && node_virial.dim() == 2 &&
+                  node_virial.size(1) == 9,
+              "frame_virial_sum: node_virial must be contiguous fp32 with "
+              "shape (N, 9)");
+  TORCH_CHECK(node_virial.is_cuda() && n_node_per_frame.is_cuda() &&
+                  node_virial.device() == n_node_per_frame.device() &&
+                  n_node_per_frame.dim() == 1,
+              "frame_virial_sum: node_virial and a one-dimensional "
+              "n_node_per_frame must share one CUDA device");
+  const long node_count = node_virial.size(0);
+  const long frame_count = n_node_per_frame.size(0);
+  auto virial = torch::zeros({frame_count, 3, 3},
+                             node_virial.options().dtype(torch::kFloat64));
+  if (node_count == 0 || frame_count == 0) {
+    return virial;
+  }
+  auto frame_row_ptr = frame_row_pointer(n_node_per_frame);
+  const int partial_count = frame_segment_partials(node_count, frame_count);
+  auto partial =
+      torch::empty({frame_count * 9, partial_count}, virial.options());
+  launch_frame_segment_sum(
+      node_count, frame_count, frame_row_ptr, node_virial.data_ptr(),
+      partial, virial.data_ptr(), at::cuda::getCurrentCUDAStream());
+  return virial;
+}
+
 // Per-frame total of a scalar carried on the node axis, the energy being the
 // only such quantity. Scattering it with an index add would serialize one
 // atomic per node on a single address whenever the batch holds one frame,
diff --git a/source/op/pt/fp32_ridge.h b/source/op/pt/fp32_ridge.h
new file mode 100644
index 0000000000..c5c4d752bf
--- /dev/null
+++ b/source/op/pt/fp32_ridge.h
@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// Ridge point of the current device, which the operators that trade memory
+// traffic for float32 SIMT arithmetic compare with their arithmetic intensity.
+
+#pragma once
+
+#include 
+
+#include 
+#include 
+
+/// Ridge point of the current device: peak fp32 FMA throughput over DRAM
+/// bandwidth, in FLOP per byte.
+///
+/// A fused kernel that keeps its intermediate tensors on chip saves memory
+/// traffic at the cost of arithmetic efficiency, so whether it pays scales
+/// with this ratio: the routing gate of the fused SO(2) convolution normalizes
+/// its measured threshold by the ridge of the card it was calibrated on, and
+/// the fused fitting network compares the ridge with the arithmetic intensity
+/// of the per-layer kernels it replaces. Both decisions thereby follow the
+/// actual part rather than an architecture name. Attribute queries cover every
+/// supported toolkit; the FMA-per-SM width is 128 lanes on every consumer and
+/// data-center part since Ampere except the A100 die (64), which the
+/// compute-capability pair identifies. Each device answers once.
+inline double fp32_ridge() {
+  static std::mutex mutex;
+  static std::map ridges;
+  int dev = 0;
+  cudaGetDevice(&dev);
+  const std::lock_guard lock(mutex);
+  const auto known = ridges.find(dev);
+  if (known != ridges.end()) {
+    return known->second;
+  }
+  int sm_count = 0, clock_khz = 0, mem_clock_khz = 0, bus_width = 0;
+  int major = 0, minor = 0;
+  cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
+  cudaDeviceGetAttribute(&clock_khz, cudaDevAttrClockRate, dev);
+  cudaDeviceGetAttribute(&mem_clock_khz, cudaDevAttrMemoryClockRate, dev);
+  cudaDeviceGetAttribute(&bus_width, cudaDevAttrGlobalMemoryBusWidth, dev);
+  cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev);
+  cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev);
+  const int lanes = (major < 8 || (major == 8 && minor == 0)) ? 64 : 128;
+  const double flops = 2.0 * lanes * sm_count * (clock_khz * 1e3);
+  const double bytes = (mem_clock_khz * 1e3) * (bus_width / 8.0) * 2.0;
+  const double ridge = bytes > 0.0 ? flops / bytes : 0.0;
+  ridges.emplace(dev, ridge);
+  return ridge;
+}
diff --git a/source/op/pt/graph_fitting.cu b/source/op/pt/graph_fitting.cu
index e0e5575fc1..dec91756da 100644
--- a/source/op/pt/graph_fitting.cu
+++ b/source/op/pt/graph_fitting.cu
@@ -6,37 +6,76 @@
 //   h_0 = act(x @ W_0 + b_0)                      (+ identity residual when
 //   h_l = act(h_{l-1} @ W_l + b_l)                  the layer is square)
 //   e   = h_{L-1} @ w_head + b_head + bias_atom_e[atype]   (fp64 output)
-// The GEMMs run on cuBLAS in pedantic fp32 (TF32 off); each layer's bias,
-// activation and residual collapse into one elementwise epilogue kernel.
+// Each layer is one FP32 SIMT GEMM kernel with the elementwise work fused into
+// it: the epilogue stores the pre-activation, adds the bias, applies the
+// activation and the identity residual, and stores the activation; the last
+// layer reduces its activation against the head weight instead, so no hidden
+// tensor is read back for the energy.
 // The backward (upstream d_e, a unit vector for the energy reduction) chains
 //   dh_{L-1} = d_e * w_head^T
 //   dpre_l   = dh_l * act'(pre_l + b_l)
 //   dh_{l-1} = dpre_l @ W_l^T  (+ dh_l identity residual)
 //   d_x      = dpre_0 @ W_0^T
-// with the elementwise steps fused likewise.
+// with dpre formed while the GEMM stages its A operand -- the seed layer forms
+// dh_{L-1} there as well, so it never exists as a tensor -- and the identity
+// residual added in the GEMM epilogue. A layer without identity residual
+// consumes only dpre, so when its GEMM spans several column tiles, each of
+// which would form the same dpre again, the epilogue of the layer above emits
+// that dpre directly and it is formed once.
+//
+// The GEMMs accumulate in fp32 on the SIMT cores with one fused multiply-add
+// per term, without TF32 or split-K reordering, so the fp32 potential-energy
+// surface is deterministic for MD. Every reduction runs in a fixed order
+// without atomics: each GEMM output sums its contraction in ascending order,
+// and the energy head reduces each 32-channel group of its dot product by a
+// fixed butterfly and sums the groups in ascending order. Neither order
+// depends on the tile shape. The tile routing also decides whether a dpre is
+// emitted or formed while staging, and both evaluate the same expression on
+// the same operands in the same order, so the results are bitwise independent
+// of the tile routing and of how a caller splits the node axis into runs.
 //
 // The saved state is the pre-activation of every layer, which each GEMM writes
-// directly. The backward re-derives the activation derivative from it, so the
-// epilogue writes only the activation and the layer costs one full-tensor
-// store less than a formulation that also materializes the derivative.
+// directly. The backward re-derives the activation derivative from it, so
+// besides it the forward stores only the activation, one full-tensor store per
+// layer less than a formulation that also materializes the derivative.
+//
+// Inference, which knows the energy cotangent before the forward runs, can
+// take a fused network instead: one kernel carries a panel of nodes through
+// every layer, the head and the backward with the activations in shared
+// memory, and reproduces the energy and the input cotangent of the per-layer
+// path bit for bit. It serves the networks of one hidden width it is compiled
+// for, on runs whose panels split over clusters of CTAs that the device keeps
+// resident at once or, as single CTAs, cover enough of the device, except
+// where the per-layer kernels are compute-bound on the device;
+// dispatch_fused_shape lists the shapes of each width. A fitting-route table
+// measured at freeze time on the deployment device names another route for
+// the run sizes at which the measurement departs from these rules
+// (fitting_route).
 //
-// All tensors here are node-scale (atoms, not edges); the fusion removes
-// kernel launches and aten glue rather than FLOPs. The head bias arrives as
-// a device tensor so that symbolic tracing never reads a value host-side.
+// All tensors here are node-scale (atoms, not edges); the fusion removes the
+// full-tensor elementwise passes rather than FLOPs. The head bias arrives as a
+// device tensor so that symbolic tracing never reads a value host-side.
 
 #include 
+#include 
+#include 
 #include 
-#include 
 #include 
 #include 
 #include 
 
 #include 
+#include 
+#include 
+#include 
+#include 
+#include 
 #include 
-#include 
+#include 
 #include 
 #include 
 
+#include "fp32_ridge.h"
 #include "graph_ops.h"
 
 namespace {
@@ -47,51 +86,6 @@ namespace {
     TORCH_CHECK(err == cudaSuccess, what, ": ", cudaGetErrorString(err)); \
   } while (0)
 
-cublasHandle_t cublas_handle() {
-  // A cuBLAS handle is device-bound and unsafe to share across threads, so
-  // cache one per device in thread-local storage. Pedantic math keeps the fp32
-  // potential-energy surface exact (no TF32, no split-K reordering) for MD.
-  thread_local std::unordered_map handles;
-  int device = 0;
-  cudaGetDevice(&device);
-  cublasHandle_t& h = handles[device];
-  if (!h) {
-    cublasCreate(&h);
-    cublasSetMathMode(h, CUBLAS_PEDANTIC_MATH);
-  }
-  return h;
-}
-
-// Row-major C(m, n) = A(m, k) @ B(k, n) + beta * C.
-void gemm_nn(cudaStream_t stream,
-             const float* a,
-             const float* b,
-             float* c,
-             int m,
-             int n,
-             int k,
-             float beta = 0.f) {
-  cublasSetStream(cublas_handle(), stream);
-  const float alpha = 1.f;
-  cublasSgemm(cublas_handle(), CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, b, n,
-              a, k, &beta, c, n);
-}
-
-// Row-major C(m, n) = A(m, k) @ B(n, k)^T + beta * C.
-void gemm_nt(cudaStream_t stream,
-             const float* a,
-             const float* b,
-             float* c,
-             int m,
-             int n,
-             int k,
-             float beta = 0.f) {
-  cublasSetStream(cublas_handle(), stream);
-  const float alpha = 1.f;
-  cublasSgemm(cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, n, m, k, &alpha, b, k,
-              a, k, &beta, c, n);
-}
-
 // sigma(z) = (1 + tanh(z/2)) / 2. The identity is not the detour it looks
 // like: CUDA's tanhf is a branch-free sequence of one MUFU.EX2 and one
 // MUFU.RCP, whereas an accurate expf carries a branchy slow path and the
@@ -112,158 +106,1779 @@ __device__ __forceinline__ float act_value(float z) {
   if constexpr (ACT == 0) {
     return tanhf(z);
   }
-  return z * sigmoid(z);
-}
+  return z * sigmoid(z);
+}
+
+template 
+__device__ __forceinline__ float act_derivative(float z) {
+  if constexpr (ACT == 0) {
+    const float a = tanhf(z);
+    return 1.f - a * a;
+  }
+  const float s = sigmoid(z);
+  return s * (1.f + z * (1.f - s));
+}
+
+long ceil_div(long a, long b) { return (a + b - 1) / b; }
+
+// ---------------------------------------------------------------------------
+// Vector helpers. Every hidden width is a positive multiple of four and every
+// weight, bias and head weight 16-byte aligned (check_weights,
+// check_parameters), so every hidden row, bias and head weight is addressed in
+// aligned float4 quads.
+// ---------------------------------------------------------------------------
+
+__device__ __forceinline__ float4 zero4() {
+  return make_float4(0.f, 0.f, 0.f, 0.f);
+}
+
+__device__ __forceinline__ float4 load4(const float* p) {
+  return __ldg(reinterpret_cast(p));
+}
+
+/// Load through L2 only, for data that the kernel itself writes while it runs,
+/// which the read-only path of ``load4`` must not serve.
+__device__ __forceinline__ float4 load4_coherent(const float* p) {
+  return __ldcg(reinterpret_cast(p));
+}
+
+__device__ __forceinline__ float4 lds4(const float* p) {
+  return *reinterpret_cast(p);
+}
+
+__device__ __forceinline__ float2 lds2(const float* p) {
+  return *reinterpret_cast(p);
+}
+
+__device__ __forceinline__ void store4(float* p, float4 v) {
+  *reinterpret_cast(p) = v;
+}
+
+__device__ __forceinline__ float4 add4(float4 a, float4 b) {
+  return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
+}
+
+/// a + b per element as explicitly rounded additions. The forward identity
+/// residual follows the activation product, which a compiler may contract
+/// with the addition into one FMA in one kernel and not in another; both
+/// forward epilogues add it with this function, so the fused network and the
+/// per-layer kernels round the same operations whatever the code generation.
+/// The backward residual adds two values that no product produces, so a plain
+/// addition rounds it identically in every kernel.
+__device__ __forceinline__ float4 add4_rn(float4 a, float4 b) {
+  return make_float4(__fadd_rn(a.x, b.x), __fadd_rn(a.y, b.y),
+                     __fadd_rn(a.z, b.z), __fadd_rn(a.w, b.w));
+}
+
+/// s * a + b per element, each with a single rounding.
+__device__ __forceinline__ float4 fma4(float s, float4 a, float4 b) {
+  return make_float4(fmaf(s, a.x, b.x), fmaf(s, a.y, b.y), fmaf(s, a.z, b.z),
+                     fmaf(s, a.w, b.w));
+}
+
+__device__ __forceinline__ float4 scale4(float s, float4 v) {
+  return make_float4(s * v.x, s * v.y, s * v.z, s * v.w);
+}
+
+template 
+__device__ __forceinline__ float4 act_value4(float4 z) {
+  return make_float4(act_value(z.x), act_value(z.y),
+                     act_value(z.z), act_value(z.w));
+}
+
+template 
+__device__ __forceinline__ float4 act_derivative4(float4 z) {
+  return make_float4(act_derivative(z.x), act_derivative(z.y),
+                     act_derivative(z.z), act_derivative(z.w));
+}
+
+/// Request the line holding ``p`` into L2.
+__device__ __forceinline__ void prefetch_l2(const float* p) {
+  asm volatile("prefetch.global.L2 [%0];" ::"l"(p));
+}
+
+/// Depth slices by which the streamed node-major operands are requested into
+/// L2 ahead of their register load. The register stage hides one slice of
+/// latency; the prefetch covers the DRAM latency beyond it, which dominates
+/// once the run outgrows L2.
+constexpr int kLookahead = 2;
+
+// ---------------------------------------------------------------------------
+// GEMM tiles.
+//
+// A CTA computes a ``rows`` x ``cols`` block of C = A @ B, A the node-major
+// (nodes, K) operand and B the (K, cols) weight operand, in ``depth``-deep
+// slices staged through double-buffered shared memory: the global loads of
+// slice s + 1 are in flight in registers while slice s is multiplied. A warp
+// is 4 x 8 lanes and a lane owns tm x tn accumulators as blocks of
+// row_group x 4, at rows r + h * row_stride + [0, row_group) and columns
+// c + g * col_stride + [0, 4) from its origin (r, c), where a row group is a
+// quad, or a pair for a lane of two rows. The four lane rows of a warp read
+// adjacent 16- or 8-byte words of a shared A row and its eight lane columns
+// adjacent 16-byte words of a shared B row. The lane rows and columns are
+// lane / 8 and lane % 8, or, for a tile with ``paired_columns``, lane / 2 % 4
+// and 2 * (lane / 8) + lane % 2, which gives each quarter-warp two adjacent
+// B words instead of all eight. On sm_120 a 16-byte B read then takes two
+// shared wavefronts instead of four, while the A reads keep theirs: one per
+// 8-byte and two per 16-byte read in either layout.
+// ---------------------------------------------------------------------------
+
+/// Shared rows carry four floats of skew. Every tile extent is a multiple of
+/// 32 floats, so the row stride is congruent to 4 modulo the 32 banks, which
+/// spreads the transposing scalar stores of a warp over 32 banks while every
+/// row stays 16-byte aligned.
+constexpr int kSkew = 4;
+
+template 
+struct GemmTile {
+  static constexpr int rows = kRows;
+  static constexpr int cols = kCols;
+  static constexpr int tm = kThreadRows;
+  static constexpr int tn = kThreadCols;
+  static constexpr int depth = kDepth;
+  static constexpr int warp_cols = kWarpCols;
+  static constexpr int threads = 32 * kWarpRows * kWarpCols;
+  static constexpr int min_blocks = kMinBlocks;
+  static constexpr int row_group = 4;
+  static constexpr bool paired_columns = false;
+  static constexpr int row_stride = kRows / (kThreadRows / 4);
+  static constexpr int col_stride = kCols / (kThreadCols / 4);
+  static_assert(kThreadRows % 4 == 0 && kThreadCols % 4 == 0,
+                "a lane owns whole row and column quads");
+  static_assert(kRows == 4 * kThreadRows * kWarpRows &&
+                    kCols == 8 * kThreadCols * kWarpCols,
+                "a warp covers 4 tm rows and 8 tn columns as 4 x 8 lanes");
+  static_assert(kRows % 32 == 0 && kCols % 32 == 0 && kDepth % 8 == 0,
+                "tile extents are whole 32-row, 32-column, 8-deep blocks");
+  static_assert(kCols / (kThreadCols / 4) % 32 == 0,
+                "the eight lane columns of a quad cover one 32-column group");
+};
+
+/// Large tile, 8 x 8 accumulators per lane, for grids of several waves.
+using LargeTile = GemmTile<64, 128, 8, 8, 8, 2, 2, 3>;
+/// Small tile with a four times finer grid for runs that would leave the
+/// device partly idle under the large one. Its 27 KB of shared memory keep
+/// three CTAs resident per SM on sm_120.
+using SmallTile = GemmTile<32, 64, 4, 4, 32, 2, 2, 3>;
+
+template 
+struct GemmStorage {
+  float a[2][Tile::depth][Tile::rows + kSkew];
+  float b[2][Tile::depth][Tile::cols + kSkew];
+};
+
+template 
+using Accumulators = float[Tile::tm][Tile::tn];
+
+/// Tile-local row and column of a lane's first accumulator.
+struct LaneOrigin {
+  int row;
+  int col;
+};
+
+/// Lane row and lane column of ``lane`` within the 4 x 8 lanes of its warp.
+template 
+__device__ __forceinline__ int lane_row(int lane) {
+  return Tile::paired_columns ? lane / 2 % 4 : lane / 8;
+}
+
+template 
+__device__ __forceinline__ int lane_col(int lane) {
+  return Tile::paired_columns ? 2 * (lane / 8) + lane % 2 : lane % 8;
+}
+
+/// Lane distance between lane columns ``d`` apart, for d of 1, 2 or 4 within
+/// one aligned group of eight, which a butterfly over the columns uses.
+template 
+__device__ __forceinline__ constexpr int lane_col_stride(int d) {
+  return Tile::paired_columns && d > 1 ? 4 * d : d;
+}
+
+template 
+__device__ __forceinline__ LaneOrigin lane_origin() {
+  const int warp = threadIdx.x / 32;
+  const int lane = threadIdx.x % 32;
+  return {
+      ((warp / Tile::warp_cols) * 4 + lane_row(lane)) * Tile::row_group,
+      ((warp % Tile::warp_cols) * 8 + lane_col(lane)) * 4};
+}
+
+/// Tile-local row of accumulator row ``i``.
+template 
+__device__ __forceinline__ int accumulator_row(const LaneOrigin& o, int i) {
+  return o.row + (i / Tile::row_group) * Tile::row_stride + i % Tile::row_group;
+}
+
+/// Tile-local column of the first accumulator of column quad ``g``.
+template 
+__device__ __forceinline__ int accumulator_col(const LaneOrigin& o, int g) {
+  return o.col + g * Tile::col_stride;
+}
+
+template 
+__device__ __forceinline__ float4
+accumulator_quad(const Accumulators& acc, int i, int g) {
+  return make_float4(acc[i][4 * g], acc[i][4 * g + 1], acc[i][4 * g + 2],
+                     acc[i][4 * g + 3]);
+}
+
+/// CTA tile of a one-dimensional grid: consecutive CTAs walk the column tiles
+/// of one row tile, so they share their A rows while those are hot in L2.
+struct TileIndex {
+  long row0;
+  int col0;
+};
+
+template 
+__device__ __forceinline__ TileIndex tile_index(int width) {
+  const int col_tiles = (width + Tile::cols - 1) / Tile::cols;
+  return {(long)(blockIdx.x / col_tiles) * Tile::rows,
+          (int)(blockIdx.x % col_tiles) * Tile::cols};
+}
+
+// ---------------------------------------------------------------------------
+// Operand stagers. ``fetch(k0)`` issues the global loads of the slice
+// [k0, k0 + Tile::depth) of the contraction axis, of length ``extent``, into
+// registers and ``commit`` writes them to shared memory in the depth-major
+// layout the fragments read. Depths past the extent are zero-filled; rows and
+// columns past the matrix are clamped to its last row or column, whose finite
+// values reach only accumulators that are never stored.
+// ---------------------------------------------------------------------------
+
+/// Thread map of a ``kLines`` x depth slice read as float4 along depth and
+/// stored depth-major. Lane pairs cover the 32-byte segment of one line in
+/// each 8-deep group of the slice, so a warp reads 16 lines and its scalar
+/// stores land on 32 distinct banks.
+template 
+struct TransposedSlice {
+  static constexpr int count = kLines * Tile::depth / 4 / Tile::threads;
+  static_assert(count >= 1 && count * 4 * Tile::threads == kLines * Tile::depth,
+                "the threads must tile the slice evenly");
+  static_assert(2 * kLines % 32 == 0, "a warp stays within one 8-deep group");
+  __device__ static int line(int s) {
+    const int i = threadIdx.x + s * Tile::threads;
+    return (i / 2) % kLines;
+  }
+  __device__ static int quad(int s) {
+    const int i = threadIdx.x + s * Tile::threads;
+    return i % 2 + 2 * (i / (2 * kLines));
+  }
+};
+
+template 
+__device__ __forceinline__ void store_transposed(float (*dst)[kLines + kSkew],
+                                                 const float4 (&v)[kCount]) {
+  using Map = TransposedSlice;
+#pragma unroll
+  for (int s = 0; s < kCount; ++s) {
+    const int q = 4 * Map::quad(s);
+    const int l = Map::line(s);
+    dst[q][l] = v[s].x;
+    dst[q + 1][l] = v[s].y;
+    dst[q + 2][l] = v[s].z;
+    dst[q + 3][l] = v[s].w;
+  }
+}
+
+/// A operand read as stored from a row-major (nodes, K) tensor: the layer
+/// input of the forward, or a dpre emitted for the backward.
+template 
+struct InputStager;
+
+/// float4 loads, for 16-byte aligned rows.
+template 
+struct InputStager {
+  using Map = TransposedSlice;
+  const float* src[Map::count];
+  float4 v[Map::count];
+  int extent;
+
+  __device__ InputStager(const float* x, long row0, long nodes, int extent_)
+      : extent(extent_) {
+#pragma unroll
+    for (int s = 0; s < Map::count; ++s) {
+      const long row = min(row0 + Map::line(s), nodes - 1);
+      src[s] = x + row * extent + 4 * Map::quad(s);
+    }
+  }
+
+  __device__ __forceinline__ void fetch(int k0) {
+#pragma unroll
+    for (int s = 0; s < Map::count; ++s) {
+      const int k = k0 + 4 * Map::quad(s);
+      v[s] = k < extent ? load4(src[s] + k0) : zero4();
+      if (k + kLookahead * Tile::depth < extent) {
+        prefetch_l2(src[s] + k0 + kLookahead * Tile::depth);
+      }
+    }
+  }
+
+  __device__ __forceinline__ void commit(float (*a)[Tile::rows + kSkew]) {
+    store_transposed(a, v);
+  }
+};
+
+/// Scalar loads, for any row stride: eight lanes cover the 32-byte segment of
+/// one line and the warp's stores land on 32 distinct banks.
+template 
+struct InputStager {
+  static constexpr int count = Tile::rows * Tile::depth / Tile::threads;
+  static_assert(count >= 1 && count * Tile::threads == Tile::rows * Tile::depth,
+                "the threads must tile the slice evenly");
+  const float* src[count];
+  float v[count];
+  int extent;
+
+  __device__ static int line(int s) {
+    const int i = threadIdx.x + s * Tile::threads;
+    return (i / 8) % Tile::rows;
+  }
+  __device__ static int col(int s) {
+    const int i = threadIdx.x + s * Tile::threads;
+    return i % 8 + 8 * (i / (8 * Tile::rows));
+  }
+
+  __device__ InputStager(const float* x, long row0, long nodes, int extent_)
+      : extent(extent_) {
+#pragma unroll
+    for (int s = 0; s < count; ++s) {
+      const long row = min(row0 + line(s), nodes - 1);
+      src[s] = x + row * extent + col(s);
+    }
+  }
+
+  __device__ __forceinline__ void fetch(int k0) {
+#pragma unroll
+    for (int s = 0; s < count; ++s) {
+      const int k = k0 + col(s);
+      v[s] = k < extent ? __ldg(src[s] + k0) : 0.f;
+      if (k + kLookahead * Tile::depth < extent) {
+        prefetch_l2(src[s] + k0 + kLookahead * Tile::depth);
+      }
+    }
+  }
+
+  __device__ __forceinline__ void commit(float (*a)[Tile::rows + kSkew]) {
+#pragma unroll
+    for (int s = 0; s < count; ++s) {
+      a[col(s)][line(s)] = v[s];
+    }
+  }
+};
+
+/// Forward B operand: the row-major (K, width) weight, read as float4
+/// along its output channels and stored as is.
+template 
+struct WeightStager {
+  static constexpr int quads = Tile::cols / 4;
+  static constexpr int count = Tile::depth * quads / Tile::threads;
+  static_assert(count >= 1 && count * Tile::threads == Tile::depth * quads &&
+                    Tile::threads % quads == 0,
+                "the threads must tile the slice evenly");
+  const float* src;
+  float4 v[count];
+  int extent;
+  int width;
+
+  __device__ static int line(int s) {
+    return (threadIdx.x + s * Tile::threads) / quads;
+  }
+
+  __device__ WeightStager(const float* w, int col0, int extent_, int width_)
+      : src(w + min(col0 + 4 * (int)(threadIdx.x % quads), width_ - 4)),
+        extent(extent_),
+        width(width_) {}
+
+  __device__ __forceinline__ void fetch(int k0) {
+#pragma unroll
+    for (int s = 0; s < count; ++s) {
+      const int k = k0 + line(s);
+      v[s] = k < extent ? load4(src + (long)k * width) : zero4();
+    }
+  }
+
+  __device__ __forceinline__ void commit(float (*b)[Tile::cols + kSkew]) {
+#pragma unroll
+    for (int s = 0; s < count; ++s) {
+      store4(&b[line(s)][4 * (threadIdx.x % quads)], v[s]);
+    }
+  }
+};
+
+/// Backward B operand: W^T for the row-major (width, K) weight W, read
+/// as float4 along W's rows (the contraction axis) and stored transposed.
+template 
+struct WeightTransposeStager {
+  using Map = TransposedSlice;
+  const float* src[Map::count];
+  float4 v[Map::count];
+  int extent;
+
+  __device__ WeightTransposeStager(const float* w,
+                                   int col0,
+                                   int extent_,
+                                   int width)
+      : extent(extent_) {
+#pragma unroll
+    for (int s = 0; s < Map::count; ++s) {
+      const int n = min(col0 + Map::line(s), width - 1);
+      src[s] = w + (long)n * extent + 4 * Map::quad(s);
+    }
+  }
+
+  __device__ __forceinline__ void fetch(int k0) {
+#pragma unroll
+    for (int s = 0; s < Map::count; ++s) {
+      v[s] = k0 + 4 * Map::quad(s) < extent ? load4(src[s] + k0) : zero4();
+    }
+  }
+
+  __device__ __forceinline__ void commit(float (*b)[Tile::cols + kSkew]) {
+    store_transposed(b, v);
+  }
+};
+
+/// Where a backward layer takes its A operand dpre = dh * act'(pre + b) from.
+enum class Cotangent {
+  kHead,    //!< Seed layer: dh = d_e[n] * w_head, dpre formed while staging.
+  kChain,   //!< dh read from the layer above, dpre formed while staging.
+  kFormed,  //!< dpre emitted by the epilogue of the layer above.
+};
+
+/// Backward A operand formed while staging, so the derivative never
+/// round-trips through memory. The seed layer forms dh = d_e[n] * w_head on the
+/// fly as well, so no head cotangent tensor exists.
+template 
+struct CotangentStager {
+  using Map = TransposedSlice;
+  const float* pre_src[Map::count];
+  const float* dh_src[Map::count];
+  float de[Map::count];
+  float4 pre[Map::count];
+  float4 g[Map::count];
+  float4 bias[Map::count];
+  const float* b;
+  const float* w_head;
+  int extent;
+
+  __device__ CotangentStager(const float* pre_,
+                             const float* dh,
+                             const double* d_e,
+                             const float* w_head_,
+                             const float* b_,
+                             long row0,
+                             long nodes,
+                             int extent_)
+      : b(b_), w_head(w_head_), extent(extent_) {
+#pragma unroll
+    for (int s = 0; s < Map::count; ++s) {
+      const long row = min(row0 + Map::line(s), nodes - 1);
+      const long offset = row * extent + 4 * Map::quad(s);
+      pre_src[s] = pre_ + offset;
+      if constexpr (kSeed) {
+        de[s] = (float)d_e[row];
+      } else {
+        dh_src[s] = dh + offset;
+      }
+    }
+  }
+
+  __device__ __forceinline__ void fetch(int k0) {
+#pragma unroll
+    for (int s = 0; s < Map::count; ++s) {
+      const int k = k0 + 4 * Map::quad(s);
+      const bool valid = k < extent;
+      pre[s] = valid ? load4(pre_src[s] + k0) : zero4();
+      if constexpr (kSeed) {
+        g[s] = valid ? load4(w_head + k) : zero4();
+      } else {
+        g[s] = valid ? load4(dh_src[s] + k0) : zero4();
+      }
+      bias[s] = valid && b ? load4(b + k) : zero4();
+      if (k + kLookahead * Tile::depth < extent) {
+        prefetch_l2(pre_src[s] + k0 + kLookahead * Tile::depth);
+        if constexpr (!kSeed) {
+          prefetch_l2(dh_src[s] + k0 + kLookahead * Tile::depth);
+        }
+      }
+    }
+  }
+
+  __device__ __forceinline__ void commit(float (*a)[Tile::rows + kSkew]) {
+    float4 d[Map::count];
+#pragma unroll
+    for (int s = 0; s < Map::count; ++s) {
+      const float4 dh = kSeed ? scale4(de[s], g[s]) : g[s];
+      const float4 slope = act_derivative4(add4(pre[s], bias[s]));
+      d[s] = make_float4(dh.x * slope.x, dh.y * slope.y, dh.z * slope.z,
+                         dh.w * slope.w);
+    }
+    store_transposed(a, d);
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Main loop. Every accumulator sums the contraction axis in ascending order
+// with one fused multiply-add per term, independent of the tile shape.
+// ---------------------------------------------------------------------------
+
+template 
+__device__ __forceinline__ void multiply_slice(
+    const float (*a)[Tile::rows + kSkew],
+    const float (*b)[Tile::cols + kSkew],
+    const LaneOrigin& o,
+    Accumulators& acc) {
+#pragma unroll
+  for (int k = 0; k < Tile::depth; ++k) {
+    float av[Tile::tm];
+    float bv[Tile::tn];
+#pragma unroll
+    for (int h = 0; h < Tile::tm / Tile::row_group; ++h) {
+      const float* row = &a[k][o.row + h * Tile::row_stride];
+      if constexpr (Tile::row_group == 4) {
+        const float4 v = lds4(row);
+        av[4 * h] = v.x;
+        av[4 * h + 1] = v.y;
+        av[4 * h + 2] = v.z;
+        av[4 * h + 3] = v.w;
+      } else {
+        const float2 v = lds2(row);
+        av[2 * h] = v.x;
+        av[2 * h + 1] = v.y;
+      }
+    }
+#pragma unroll
+    for (int g = 0; g < Tile::tn / 4; ++g) {
+      const float4 v = lds4(&b[k][o.col + g * Tile::col_stride]);
+      bv[4 * g] = v.x;
+      bv[4 * g + 1] = v.y;
+      bv[4 * g + 2] = v.z;
+      bv[4 * g + 3] = v.w;
+    }
+#pragma unroll
+    for (int i = 0; i < Tile::tm; ++i) {
+#pragma unroll
+      for (int j = 0; j < Tile::tn; ++j) {
+        acc[i][j] = fmaf(av[i], bv[j], acc[i][j]);
+      }
+    }
+  }
+}
+
+template 
+__device__ __forceinline__ void zero_accumulators(Accumulators& acc) {
+#pragma unroll
+  for (int i = 0; i < S::tm; ++i) {
+#pragma unroll
+    for (int j = 0; j < S::tn; ++j) {
+      acc[i][j] = 0.f;
+    }
+  }
+}
+
+template 
+__device__ __forceinline__ void gemm_mainloop(GemmStorage& smem,
+                                              StagerA& sa,
+                                              StagerB& sb,
+                                              int extent,
+                                              const LaneOrigin& o,
+                                              Accumulators& acc) {
+  zero_accumulators(acc);
+  const int slices = (extent + Tile::depth - 1) / Tile::depth;
+  sa.fetch(0);
+  sb.fetch(0);
+  sa.commit(smem.a[0]);
+  sb.commit(smem.b[0]);
+  __syncthreads();
+  for (int s = 0; s < slices; ++s) {
+    const int buf = s & 1;
+    sa.fetch((s + 1) * Tile::depth);
+    sb.fetch((s + 1) * Tile::depth);
+    multiply_slice(smem.a[buf], smem.b[buf], o, acc);
+    sa.commit(smem.a[buf ^ 1]);
+    sb.commit(smem.b[buf ^ 1]);
+    __syncthreads();
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Forward layer: pre = x @ W, y = act(pre + b) (+ x).
+//
+// The last layer stores no activation. Its head dot products are reduced in
+// 32-channel groups: each lane takes the dot of its own four channels, and
+// the eight lanes of a group combine by a butterfly, which yields the same
+// value on every lane and in every tile shape. The groups are then summed in
+// ascending order, inside the CTA when it spans the whole width and otherwise
+// by head_finalize_kernel from the per-group partials.
+// ---------------------------------------------------------------------------
+
+constexpr int kHeadGroup = 32;
+
+struct LayerForward {
+  const float* x;             //!< (nodes, din) layer input.
+  const float* w;             //!< (din, dout) weight.
+  const float* b;             //!< (dout,) bias or nullptr.
+  float* pre;                 //!< (nodes, dout) saved pre-activation.
+  float* y;                   //!< (nodes, dout) activation, null for the last.
+  const float* w_head;        //!< (dout,) head weight.
+  const float* b_head;        //!< (1,) head bias or nullptr.
+  const double* bias_atom_e;  //!< (ntypes,) per-type energy bias.
+  const long* atype;          //!< (nodes,) atom types.
+  float* head_partial;        //!< (groups, nodes) partial dots of the head.
+  double* e;                  //!< (nodes,) atomic energy.
+  long nodes;
+  int din;
+  int dout;
+  int residual;
+};
+
+/// e = (sum of the group partials in ascending order + b_head) +
+/// bias_atom_e[atype], in fp64 after the fp32 head.
+__device__ __forceinline__ double head_energy(const float* partial,
+                                              long stride,
+                                              int groups,
+                                              const float* b_head,
+                                              double bias) {
+  float acc = 0.f;
+  for (int c = 0; c < groups; ++c) {
+    acc += partial[c * stride];
+  }
+  return (double)(acc + (b_head ? b_head[0] : 0.f)) + bias;
+}
+
+template 
+__global__ void __launch_bounds__(Tile::threads, Tile::min_blocks)
+    layer_forward_kernel(const LayerForward p) {
+  __shared__ __align__(16) GemmStorage smem;
+  const LaneOrigin o = lane_origin();
+  const TileIndex tile = tile_index(p.dout);
+  Accumulators acc;
+  {
+    InputStager sa(p.x, tile.row0, p.nodes, p.din);
+    WeightStager sb(p.w, tile.col0, p.din, p.dout);
+    gemm_mainloop(smem, sa, sb, p.din, o, acc);
+  }
+
+  // === Step 1. Pre-activation, activation and the lane head dots ===
+  const bool last = p.y == nullptr;
+  float head[Tile::tm][Tile::tn / 4];
+#pragma unroll
+  for (int i = 0; i < Tile::tm; ++i) {
+    const long m = tile.row0 + accumulator_row(o, i);
+#pragma unroll
+    for (int g = 0; g < Tile::tn / 4; ++g) {
+      const int n = tile.col0 + accumulator_col(o, g);
+      head[i][g] = 0.f;
+      if (m >= p.nodes || n >= p.dout) {
+        continue;
+      }
+      const long t = m * p.dout + n;
+      const float4 z = accumulator_quad(acc, i, g);
+      store4(p.pre + t, z);
+      float4 v = act_value4(add4(z, p.b ? load4(p.b + n) : zero4()));
+      if (p.residual) {
+        // A square layer reads its input row as aligned quads unless the
+        // input itself is misaligned.
+        const float4 r =
+            kVector ? load4(p.x + t)
+                    : make_float4(__ldg(p.x + t), __ldg(p.x + t + 1),
+                                  __ldg(p.x + t + 2), __ldg(p.x + t + 3));
+        v = add4_rn(v, r);
+      }
+      if (!last) {
+        store4(p.y + t, v);
+      } else {
+        const float4 wh = load4(p.w_head + n);
+        float q = v.x * wh.x;
+        q = fmaf(v.y, wh.y, q);
+        q = fmaf(v.z, wh.z, q);
+        head[i][g] = fmaf(v.w, wh.w, q);
+      }
+    }
+  }
+  if (!last) {
+    return;
+  }
+
+  // === Step 2. 32-channel group partials ===
+  const int groups = (p.dout + kHeadGroup - 1) / kHeadGroup;
+  const bool whole = p.dout <= Tile::cols;
+  constexpr int kTileGroups = Tile::cols / kHeadGroup;
+  float* partial = &smem.a[0][0][0];  // (kTileGroups, rows), whole tiles
+#pragma unroll
+  for (int i = 0; i < Tile::tm; ++i) {
+    const int r = accumulator_row(o, i);
+#pragma unroll
+    for (int g = 0; g < Tile::tn / 4; ++g) {
+      float q = head[i][g];
+      q += __shfl_xor_sync(0xffffffffu, q, lane_col_stride(1));
+      q += __shfl_xor_sync(0xffffffffu, q, lane_col_stride(2));
+      q += __shfl_xor_sync(0xffffffffu, q, lane_col_stride(4));
+      const int group = (tile.col0 + accumulator_col(o, g)) / kHeadGroup;
+      if (lane_col(threadIdx.x % 32) != 0 || group >= groups) {
+        continue;
+      }
+      if (whole) {
+        partial[group * Tile::rows + r] = q;
+      } else if (tile.row0 + r < p.nodes) {
+        p.head_partial[group * p.nodes + tile.row0 + r] = q;
+      }
+    }
+  }
+  if (!whole) {
+    return;
+  }
+
+  // === Step 3. Energy of a CTA spanning the whole width ===
+  static_assert(
+      kTileGroups * Tile::rows <= sizeof(GemmStorage::a) / sizeof(float),
+      "the group partials fit the A buffers");
+  __syncthreads();
+  for (int r = threadIdx.x; r < Tile::rows; r += Tile::threads) {
+    const long m = tile.row0 + r;
+    if (m < p.nodes) {
+      p.e[m] = head_energy(partial + r, Tile::rows, groups, p.b_head,
+                           p.bias_atom_e[p.atype[m]]);
+    }
+  }
+}
+
+/// Energy of a head whose groups span several column tiles.
+__global__ void head_finalize_kernel(long nodes,
+                                     int groups,
+                                     const float* __restrict__ partial,
+                                     const float* __restrict__ b_head,
+                                     const double* __restrict__ bias_atom_e,
+                                     const long* __restrict__ atype,
+                                     double* __restrict__ e) {
+  const long n = blockIdx.x * (long)blockDim.x + threadIdx.x;
+  if (n >= nodes) {
+    return;
+  }
+  e[n] = head_energy(partial + n, nodes, groups, b_head, bias_atom_e[atype[n]]);
+}
+
+// ---------------------------------------------------------------------------
+// Backward layer: dh_in = dpre @ W^T (+ dh) with dpre = dh * act'(pre + b),
+// where dpre is formed while staging or read as emitted by the layer above.
+// The layer stores dh_in or, when it emits for the layer below, that layer's
+// dpre dh_in * act'(below_pre + below_b).
+// ---------------------------------------------------------------------------
+
+struct LayerBackward {
+  const float* dh;         //!< (nodes, dout) output cotangent (kChain) or the
+                           //!< layer's dpre (kFormed); unused by the seed.
+  const double* d_e;       //!< (nodes,) energy cotangent.
+  const float* w_head;     //!< (dout,) head weight.
+  const float* pre;        //!< (nodes, dout) saved pre-activation.
+  const float* b;          //!< (dout,) bias or nullptr.
+  const float* w;          //!< (din, dout) weight.
+  float* out;              //!< (nodes, din) cotangent of the layer input, or
+                           //!< the dpre of the layer below if below_pre is set.
+  const float* below_pre;  //!< (nodes, din) pre-activation of the layer below
+                           //!< to emit dpre for, or nullptr.
+  const float* below_b;    //!< (din,) bias of the layer below or nullptr.
+  long nodes;
+  int din;
+  int dout;
+  int residual;
+};
+
+template 
+__global__ void __launch_bounds__(Tile::threads, Tile::min_blocks)
+    layer_backward_kernel(const LayerBackward p) {
+  constexpr bool kSeed = kSource == Cotangent::kHead;
+  __shared__ __align__(16) GemmStorage smem;
+  const LaneOrigin o = lane_origin();
+  const TileIndex tile = tile_index(p.din);
+  Accumulators acc;
+  {
+    using StagerA = std::conditional_t,
+                                       CotangentStager>;
+    StagerA sa = [&] {
+      if constexpr (kSource == Cotangent::kFormed) {
+        return StagerA(p.dh, tile.row0, p.nodes, p.dout);
+      } else {
+        return StagerA(p.pre, p.dh, p.d_e, p.w_head, p.b, tile.row0, p.nodes,
+                       p.dout);
+      }
+    }();
+    WeightTransposeStager sb(p.w, tile.col0, p.dout, p.din);
+    gemm_mainloop(smem, sa, sb, p.dout, o, acc);
+  }
+
+  // Rows of a width that is not a multiple of four, such as the d_x of an
+  // unpadded descriptor, and a misaligned d_x are stored per element; hidden
+  // rows are whole aligned quads. A dpre emitted for the layer below is
+  // dh_in * act'(below_pre + below_b), rounded exactly as the stager of that
+  // layer would form it from a stored dh_in.
+  const bool vector =
+      p.din % 4 == 0 && reinterpret_cast(p.out) % 16 == 0;
+#pragma unroll
+  for (int i = 0; i < Tile::tm; ++i) {
+    const long m = tile.row0 + accumulator_row(o, i);
+    if (m >= p.nodes) {
+      continue;
+    }
+    const float de = kSeed && p.residual ? (float)p.d_e[m] : 0.f;
+#pragma unroll
+    for (int g = 0; g < Tile::tn / 4; ++g) {
+      const int n = tile.col0 + accumulator_col(o, g);
+      if (n >= p.din) {
+        continue;
+      }
+      const long t = m * p.din + n;
+      float4 v = accumulator_quad(acc, i, g);
+      // Emission targets only layers without identity residual, so a layer
+      // that reads an emitted dpre has no residual to add.
+      if constexpr (kSource != Cotangent::kFormed) {
+        if (p.residual) {
+          v = kSeed ? fma4(de, load4(p.w_head + n), v)
+                    : add4(v, load4(p.dh + t));
+        }
+      }
+      if (p.below_pre) {
+        const float4 bias = p.below_b ? load4(p.below_b + n) : zero4();
+        const float4 slope =
+            act_derivative4(add4(load4(p.below_pre + t), bias));
+        v = make_float4(v.x * slope.x, v.y * slope.y, v.z * slope.z,
+                        v.w * slope.w);
+      }
+      if (vector) {
+        store4(p.out + t, v);
+      } else {
+        const float c[4] = {v.x, v.y, v.z, v.w};
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+          if (n + j < p.din) {
+            p.out[t + j] = c[j];
+          }
+        }
+      }
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Fused network. One CTA carries a panel of ``rows`` nodes through the whole
+// network and back in a single launch: every hidden layer, the energy head,
+// and the cotangent of the input. The panel's activations stay in shared
+// memory, feature-major so that they are the A operand of the next GEMM as
+// stored, and each layer replaces its input in place once every lane has read
+// it. The weights of every GEMM, forward and backward in the order they are
+// multiplied, form one stream of depth slices that flows through a ring filled
+// by asynchronous copies ahead of the slice being multiplied, across the layer
+// boundaries, so no GEMM waits for its first slices; the backward reads the
+// transposed weights, so both directions copy rows as stored. A shape that
+// serves many panels per multiprocessor keeps a short ring and lets the other
+// panels hide the latency of the stream; a shape for a run of fewer panels
+// than multiprocessors keeps a deep one, which then hides it alone. The
+// backward reads the pre-activations of every layer but the last,
+// which the forward stores; the last layer's derivative and the head
+// cotangent are formed from its accumulators. The output cotangent of a
+// layer, which the identity residual of the layer below adds, replaces the
+// pre-activation sheet the backward has just consumed.
+//
+// A run of no more panels than multiprocessors splits each panel over a
+// cluster of CTAs when the device keeps a cluster per panel resident at once:
+// CTA ``rank`` computes the columns
+// [rank * cols, (rank + 1) * cols) of every GEMM and streams only those
+// columns of each weight, while it keeps the whole activation panel, which
+// the next GEMM contracts over. After each epilogue the CTAs assemble the
+// next A operand from their column slices through an exchange sheet in global
+// memory: each CTA stores its slice there and, behind a cluster barrier, the
+// next GEMM loads the slices of the others into its shared rows, each while
+// the column slice before it is multiplied. The slices travel through L2,
+// the fastest path between the CTAs of a cluster on sm_120, which delivers 16
+// to 21 B/clk to each multiprocessor. Every column a CTA owns -- its saved
+// pre-activations and its output cotangents -- stays with that CTA through
+// the whole pass, so between GEMMs only the A operands move; the head
+// partials of its groups move once, into the first CTA, which forms the
+// energy.
+//
+// Every GEMM sums its contraction in ascending order with one fused
+// multiply-add per term, the head reduces by the 32-channel butterfly and the
+// ascending group sum of the per-layer kernels, and every elementwise step
+// evaluates the same expression on the same operands, so the energy and the
+// input cotangent are bitwise those of the per-layer path for every cluster
+// size.
+// ---------------------------------------------------------------------------
+
+/// Hidden layers the fused kernel carries.
+constexpr int kFusedLayers = 8;
+
+template 
+struct FusedShape {
+  static constexpr int width = kWidth;
+  /// CTAs that split the columns of each panel between them.
+  static constexpr int cluster = kCluster;
+  /// Quarters of the multiprocessors that the 32-node panels of a run must
+  /// cover before a single CTA per panel serves it; a smaller run leaves the
+  /// device idle enough that the per-layer kernels, which spread every layer
+  /// over it, finish first.
+  static constexpr int serve_quarters = kServeQuarters;
+  static constexpr int min_blocks = kMinBlocks;
+  /// Weight slices in the ring: the one being multiplied and those ahead.
+  static constexpr int stages = kStages;
+  static constexpr int rows = 4 * kThreadRows * kWarpRows;
+  /// Columns of every GEMM that one CTA computes.
+  static constexpr int cols = 8 * kThreadCols * kWarpCols;
+  static constexpr int tm = kThreadRows;
+  static constexpr int tn = kThreadCols;
+  /// Contraction rows of a weight slice.
+  static constexpr int depth = kDepth;
+  static constexpr int warp_cols = kWarpCols;
+  static constexpr int threads = 32 * kWarpRows * kWarpCols;
+  /// Consecutive rows of a lane: a quad, or a pair for a lane of two rows,
+  /// which gives a narrow CTA twice the lanes for the same columns.
+  static constexpr int row_group = kThreadRows < 4 ? kThreadRows : 4;
+  static constexpr bool paired_columns = true;
+  static constexpr int row_stride = rows / (kThreadRows / row_group);
+  static constexpr int col_stride = cols / (kThreadCols / 4);
+  static constexpr int act_stride = rows + kSkew;
+  static constexpr int weight_stride = cols + kSkew;
+  static constexpr int act_floats = width * act_stride;
+  static constexpr int weight_floats = kStages * depth * weight_stride;
+  static constexpr int input_floats = 2 * depth * act_stride;
+  static constexpr size_t bytes =
+      sizeof(float) * (act_floats + weight_floats + input_floats);
+  static_assert(cols * kCluster == width,
+                "the CTAs of a cluster split the hidden width");
+  static_assert(cols % kHeadGroup == 0,
+                "every CTA of a cluster owns whole head groups");
+  static_assert((kThreadRows == 2 || kThreadRows % 4 == 0) &&
+                    kThreadCols % 4 == 0,
+                "a lane owns whole row groups and column quads");
+  static_assert(rows % 32 == 0 && width % 32 == 0,
+                "the panel and the head groups are whole 32-blocks");
+  static_assert(kStages >= 2, "the ring holds at least one slice ahead");
+  static_assert(kDepth % 8 == 0 && width % kDepth == 0,
+                "the hidden contraction is whole slices");
+  static_assert((width / 32) * rows <= input_floats,
+                "the head partials fit the input staging");
+  // Every shape fits the 99 KB of shared memory that compute capability 8.6,
+  // 8.9 and 12.x let one CTA opt in to, and so every device of compute
+  // capability 8.0 or later. A device that grants less, 96 KB on 7.0 and 64 KB
+  // on 7.5, takes only the shapes it admits (launch_first_fit).
+  static_assert(bytes <= 101376,
+                "a shape fits the opt-in shared memory of 8.6, 8.9 and 12.x");
+  /// The template parameters packed into one positive integer, which names
+  /// the shape in a fitting-route table (fitting_route).
+  static constexpr int64_t key =
+      static_cast(kWidth) | static_cast(kCluster) << 9 |
+      static_cast(kThreadRows) << 15 |
+      static_cast(kThreadCols) << 21 |
+      static_cast(kWarpRows) << 27 |
+      static_cast(kWarpCols) << 33 |
+      static_cast(kMinBlocks) << 39 |
+      static_cast(kStages) << 45 | static_cast(kDepth) << 51 |
+      static_cast(kServeQuarters) << 58;
+  static_assert(kWidth < 512 && kCluster < 64 && kThreadRows < 64 &&
+                    kThreadCols < 64 && kWarpRows < 64 && kWarpCols < 64 &&
+                    kMinBlocks < 64 && kStages < 64 && kDepth < 128 &&
+                    kServeQuarters < 16,
+                "every template parameter fits its field of the key");
+};
+
+struct FusedNetwork {
+  const float* x;                 //!< (nodes, din) network input.
+  const long* atype;              //!< (nodes,) atom types.
+  const float* w[kFusedLayers];   //!< (din_l, width) weights.
+  const float* wt[kFusedLayers];  //!< (width, pitch_l) transposed weights.
+  int pitch[kFusedLayers];        //!< Row pitch of wt, din_l rounded to 4.
+  const float* b[kFusedLayers];   //!< (width,) biases or nullptr.
+  int residual[kFusedLayers];     //!< Identity residual of each layer.
+  const float* w_head;            //!< (width,) head weight.
+  const float* b_head;            //!< (1,) head bias or nullptr.
+  const double* bias_atom_e;      //!< (ntypes,) per-type energy bias.
+  const double* d_e;              //!< (nodes,) energy cotangent.
+  float* saved;                   //!< n_layer sheets of (nodes, width): the
+                                  //!< pre-activations of every layer but the
+                                  //!< last, then the exchange sheet.
+  double* e;                      //!< (nodes,) atomic energy.
+  float* d_x;                     //!< (nodes, din) input cotangent, which
+                                  //!< may take the place of ``x``.
+  long nodes;
+  int din;
+  int n_layer;
+};
+
+/// Copy 16 bytes into shared memory, zero-filled when the source is outside
+/// the operand. In an image compiled for compute capability 8.0 or later the
+/// copy is asynchronous and completes at the matching ``copy_async_wait``; an
+/// image compiled for compute capability 7.x copies at once, which the
+/// barriers of the ring order the same way.
+__device__ __forceinline__ void copy_async16(void* dst,
+                                             const void* src,
+                                             bool valid) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+  const unsigned address = static_cast(__cvta_generic_to_shared(dst));
+  asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;\n" ::"r"(address),
+               "l"(src), "r"(valid ? 16 : 0));
+#else
+  *static_cast(dst) =
+      valid ? *static_cast(src) : zero4();
+#endif
+}
+
+__device__ __forceinline__ void copy_async_commit() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+  asm volatile("cp.async.commit_group;\n" ::);
+#endif
+}
+
+template 
+__device__ __forceinline__ void copy_async_wait() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+  asm volatile("cp.async.wait_group %0;\n" ::"n"(kPending));
+#endif
+}
+
+/// A operand of the first layer: depth columns of the row-major (nodes,
+/// extent) input for the panel's rows, read per element, eight lanes per
+/// 32-byte row segment, and stored feature-major. Rows past the run are
+/// clamped to its last row and depths past the extent are zero. The input is
+/// read through L2, since the input cotangent the kernel writes may take its
+/// place (see load4_coherent).
+template 
+struct FusedInputColumns {
+  static constexpr int items = S::rows * S::depth;
+  static constexpr int count = (items + S::threads - 1) / S::threads;
+  const float* x;
+  long row0;
+  long nodes;
+  int extent;
+  float v[count];
+
+  __device__ FusedInputColumns(const float* x_,
+                               long row0_,
+                               long nodes_,
+                               int extent_)
+      : x(x_), row0(row0_), nodes(nodes_), extent(extent_) {}
+
+  __device__ __forceinline__ void fetch(int k0) {
+#pragma unroll
+    for (int s = 0; s < count; ++s) {
+      const int item = threadIdx.x + s * S::threads;
+      const long row = min(row0 + item / S::depth, nodes - 1);
+      const int k = k0 + item % S::depth;
+      v[s] = item < items && k < extent ? __ldcg(x + row * extent + k) : 0.f;
+    }
+  }
+
+  __device__ __forceinline__ void commit(float (*a)[S::act_stride]) {
+#pragma unroll
+    for (int s = 0; s < count; ++s) {
+      const int item = threadIdx.x + s * S::threads;
+      if (item < items) {
+        a[item % S::depth][item / S::depth] = v[s];
+      }
+    }
+  }
+};
+
+/// B operand of one GEMM of the fused network: rows [0, rows_valid) and the
+/// columns from ``col0`` below ``cols_valid`` of a row-major weight at row
+/// ``pitch``, in ``slices`` depth slices.
+struct FusedOperand {
+  const float* base;
+  int pitch;
+  int col0;
+  int cols_valid;
+  int rows_valid;
+  int slices;
+};
+
+/// GEMMs of the fused network in the order the kernel multiplies them: the
+/// forward layers from the first, then the backward from the last layer down,
+/// where the first layer takes one pass per block of ``width`` input columns.
+__device__ __forceinline__ int fused_gemm_count(const FusedNetwork& p,
+                                                int width) {
+  return 2 * p.n_layer - 1 + (p.din + width - 1) / width;
+}
+
+/// Operand of GEMM ``gemm`` in the order of ``fused_gemm_count``, restricted
+/// to the columns from ``col0`` of every block of ``width`` output columns.
+template 
+__device__ __forceinline__ FusedOperand fused_operand(const FusedNetwork& p,
+                                                      int gemm,
+                                                      int col0) {
+  const int last = p.n_layer - 1;
+  if (gemm <= last) {
+    const int din = gemm == 0 ? p.din : S::width;
+    return {p.w[gemm], S::width, col0,
+            S::width,  din,      (din + S::depth - 1) / S::depth};
+  }
+  const int back = gemm - last - 1;
+  const int l = back < last ? last - back : 0;
+  const int pass = back < last ? 0 : back - last;
+  return {p.wt[l],    p.pitch[l], pass * S::width + col0,
+          p.pitch[l], S::width,   S::width / S::depth};
+}
+
+/// The weight slices of every GEMM of the network as one stream through the
+/// ring, each restricted to the columns of the CTA from ``col0``. Every lane
+/// walks the same sequence: the issuing cursor runs ``stages - 1`` slices
+/// ahead of the consuming one, and crosses into the next GEMM as soon as the
+/// current one is fully issued, so the slices of a GEMM are in flight while
+/// the epilogue of the previous one runs. A lane copies the same 16-byte
+/// words of every slice, ``depth`` operand rows further down each time, so it
+/// keeps a source offset per word and derives the offsets anew only when the
+/// stream enters the next GEMM. Rows and columns outside the operand arrive
+/// as zeros.
+template 
+class WeightStream {
+ public:
+  using Ring = float (*)[S::depth][S::weight_stride];
+
+  __device__ WeightStream(const FusedNetwork& p, Ring ring, int col0)
+      : p_(p), ring_(ring), total_(fused_gemm_count(p, S::width)), col0_(col0) {
+    aim();
+#pragma unroll
+    for (int s = 0; s < S::stages - 1; ++s) {
+      issue(s);
+    }
+  }
+
+  /// Wait for the next slice, issue the slice that the ring frees in its
+  /// place, and return the ring row of the waited slice. The barrier also
+  /// releases the slot of the previous slice, which every lane has
+  /// multiplied by then.
+ __device__ __forceinline__ const float (*next()) [S::weight_stride] {
+   copy_async_wait();
+   __syncthreads();
+   issue((consumed_ + S::stages - 1) % S::stages);
+   return ring_[consumed_++ % S::stages];
+ }
+
+ private : static constexpr int kQuads = S::cols / 4;
+  static constexpr int kItems = S::depth * kQuads;
+  static_assert(kItems % S::threads == 0,
+                "the lanes copy the 16-byte words of a slice in whole rounds");
+  static constexpr int kWords = kItems / S::threads;
+
+  /// Row and column of word ``w`` of a slice that the lane copies.
+  __device__ static int line(int w) {
+    return (threadIdx.x + w * S::threads) / kQuads;
+  }
+  __device__ static int column(int w) {
+    return 4 * ((threadIdx.x + w * S::threads) % kQuads);
+  }
+
+  /// Point the offsets at the first slice of the GEMM under the cursor.
+  __device__ __forceinline__ void aim() {
+    operand_ = fused_operand(p_, gemm_, col0_);
+    slice_ = 0;
+#pragma unroll
+    for (int w = 0; w < kWords; ++w) {
+      offset_[w] = line(w) * operand_.pitch + operand_.col0 + column(w);
+    }
+  }
+
+  /// Issue the slice under the issuing cursor into ``slot``, or nothing past
+  /// the end of the stream, and close the copy group either way, so that the
+  /// group count of every lane stays aligned with the slices.
+  __device__ __forceinline__ void issue(int slot) {
+    if (gemm_ < total_) {
+      const int k0 = slice_ * S::depth;
+#pragma unroll
+      for (int w = 0; w < kWords; ++w) {
+        const bool valid = operand_.col0 + column(w) < operand_.cols_valid &&
+                           k0 + line(w) < operand_.rows_valid;
+        copy_async16(&ring_[slot][line(w)][column(w)],
+                     operand_.base + (valid ? offset_[w] : 0), valid);
+        offset_[w] += S::depth * operand_.pitch;
+      }
+      if (++slice_ == operand_.slices && ++gemm_ < total_) {
+        aim();
+      }
+    }
+    copy_async_commit();
+  }
+
+  const FusedNetwork& p_;
+  Ring ring_;
+  int total_;
+  int col0_;
+  FusedOperand operand_;  // of the GEMM under the issuing cursor
+  int gemm_ = 0;
+  int slice_ = 0;
+  int consumed_ = 0;
+  int offset_[kWords];  // of each word of the lane in the next slice
+};
+
+/// C = A @ B over ``slices`` depth slices of an A operand resident in shared
+/// memory, feature-major, and the next GEMM of the stream. The closing
+/// barrier releases the A operand to the caller's epilogue.
+template 
+__device__ __forceinline__ void fused_resident_gemm(
+    const float (*a)[S::act_stride],
+    WeightStream& stream,
+    int slices,
+    const LaneOrigin& o,
+    Accumulators& acc) {
+  zero_accumulators(acc);
+  for (int s = 0; s < slices; ++s) {
+    multiply_slice(a + s * S::depth, stream.next(), o, acc);
+  }
+  __syncthreads();
+}
+
+/// Store quad ``v`` of feature-major shared rows at (feature n, row r).
+template 
+__device__ __forceinline__ void store_feature_quad(float (*act)[S::act_stride],
+                                                   int n,
+                                                   int r,
+                                                   float4 v) {
+  act[n][r] = v.x;
+  act[n + 1][r] = v.y;
+  act[n + 2][r] = v.z;
+  act[n + 3][r] = v.w;
+}
+
+template 
+__device__ __forceinline__ float4
+load_feature_quad(const float (*act)[S::act_stride], int n, int r) {
+  return make_float4(act[n][r], act[n + 1][r], act[n + 2][r], act[n + 3][r]);
+}
+
+/// Arrive on the barrier of the cluster, with release semantics: the memory
+/// operations of the thread before it are visible to every thread of the
+/// cluster once that thread's matching ``cluster_wait`` returns. Code compiled
+/// for a virtual architecture below compute_90 has no clusters, and the
+/// dispatch selects a cluster shape only for a kernel image compiled for
+/// compute_90 or later. The cluster primitives are the compute capability 9.0
+/// intrinsics of the CUDA runtime (sm_90_rt.h).
+__device__ __forceinline__ void cluster_arrive() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  __cluster_barrier_arrive();
+#else
+  __trap();
+#endif
+}
+
+/// Wait, with acquire semantics, until every thread of the cluster has
+/// arrived on the barrier phase of this thread's last ``cluster_arrive``.
+__device__ __forceinline__ void cluster_wait() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  __cluster_barrier_wait();
+#else
+  __trap();
+#endif
+}
+
+/// ``p`` in the shared memory of CTA ``rank`` of the cluster.
+__device__ __forceinline__ float* cluster_shared(float* p, int rank) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  return static_cast(__cluster_map_shared_rank(p, rank));
+#else
+  __trap();
+  return p;
+#endif
+}
+
+/// The column slices of the A operands that the CTAs of a cluster exchange
+/// between GEMMs, through the panel's rows of the (nodes, width) exchange
+/// sheet, and the head partials that the first CTA collects. A single CTA
+/// holds every column itself, and every step is empty.
+///
+/// One sheet serves every exchange of the pass, so no CTA may store its slice
+/// of an exchange while another still loads the previous one. The cluster
+/// barrier therefore alternates two phases per exchange. A CTA arrives on the
+/// first once it has stored its slice and waits on it before it loads the
+/// slices of the others, which orders every store of the exchange before every
+/// load. It arrives on the second once its last load has landed in its shared
+/// rows and waits on it only before it stores its slice of the next exchange,
+/// which orders every load before the next stores. The pass opens with an
+/// arrival on the second phase and closes with a wait on it, so no CTA exits
+/// while the first one may still read its shared memory.
+template 
+class PanelExchange {
+  /// Feature quads of a slice per lane.
+  static constexpr int kQuads = S::cols * S::rows / 4 / S::threads;
+  static_assert(kQuads * 4 * S::threads == S::cols * S::rows &&
+                    S::rows % 16 == 0 && S::cols % 8 == 0,
+                "the lanes split a slice evenly");
+
+ public:
+  /// The quads of one slice that a lane carries from the sheet.
+  using Quads = float4[kQuads];
+
+  __device__ PanelExchange(float* sheet, long row0, int rows, int rank)
+      : sheet_(sheet + row0 * S::width), rows_(rows), rank_(rank) {
+    release();
+  }
+
+  /// Wait until every CTA has loaded the previous exchange out of the sheet.
+  __device__ __forceinline__ void acquire() const {
+    if constexpr (S::cluster > 1) {
+      cluster_wait();
+    }
+  }
+
+  /// Store quad ``v`` of this CTA's slice at (row r, columns [n, n + 4)).
+  __device__ __forceinline__ void publish(int r, int n, float4 v) const {
+    if constexpr (S::cluster > 1) {
+      if (r < rows_) {
+        store4(sheet_ + (long)r * S::width + n, v);
+      }
+    }
+  }
+
+  /// Wait until every CTA of the cluster has stored its slice.
+  __device__ __forceinline__ void open() const {
+    if constexpr (S::cluster > 1) {
+      cluster_arrive();
+      cluster_wait();
+    }
+  }
+
+  /// The first peer after ``peer`` in ascending rank, skipping this CTA, or
+  /// S::cluster past the last.
+  __device__ __forceinline__ int after(int peer) const {
+    ++peer;
+    return peer == rank_ ? peer + 1 : peer;
+  }
+
+  /// Load the lane's quads of the slice of CTA ``peer``; rows past the run
+  /// are zeros. Two lanes cover the 32-byte segment of one row.
+  __device__ __forceinline__ void load(int peer, Quads& v) const {
+#pragma unroll
+    for (int s = 0; s < kQuads; ++s) {
+      v[s] = row(s) < rows_ ? load4_coherent(sheet_ + (long)row(s) * S::width +
+                                             peer * S::cols + feature(s))
+                            : zero4();
+    }
+  }
+
+  /// Store the lane's quads of the slice of CTA ``peer`` feature-major. A
+  /// warp stores two quads of 16 rows, which lands the four scalar stores of
+  /// each quad on 32 distinct banks.
+  __device__ __forceinline__ void store(float (*act)[S::act_stride],
+                                        int peer,
+                                        const Quads& v) const {
+#pragma unroll
+    for (int s = 0; s < kQuads; ++s) {
+      store_feature_quad(act, peer * S::cols + feature(s), row(s), v[s]);
+    }
+  }
+
+  /// In the first CTA, read the (width / 32, rows) head partials of the
+  /// groups that the other CTAs own, which the barrier of the seed exchange
+  /// has made visible, into its own partial rows.
+  __device__ __forceinline__ void collect(float* partial) const {
+    if constexpr (S::cluster > 1) {
+      constexpr int kOwn = S::cols / kHeadGroup * S::rows;
+      constexpr int kAll = S::width / kHeadGroup * S::rows;
+      if (rank_ == 0) {
+        for (int item = kOwn + threadIdx.x; item < kAll; item += S::threads) {
+          partial[item] = cluster_shared(partial, item / kOwn)[item];
+        }
+      }
+    }
+  }
 
-template 
-__device__ __forceinline__ float act_derivative(float z) {
-  if constexpr (ACT == 0) {
-    const float a = tanhf(z);
-    return 1.f - a * a;
+  /// Release the sheet to the next exchange.
+  __device__ __forceinline__ void release() const {
+    if constexpr (S::cluster > 1) {
+      cluster_arrive();
+    }
   }
-  const float s = sigmoid(z);
-  return s * (1.f + z * (1.f - s));
-}
 
-long ceil_div(long a, long b) { return (a + b - 1) / b; }
+ private:
+  /// Row and first feature of the lane's quad ``s`` within a slice.
+  __device__ static int row(int s) {
+    return (threadIdx.x + s * S::threads) / 2 % S::rows;
+  }
+  __device__ static int feature(int s) {
+    const int q = threadIdx.x + s * S::threads;
+    return 4 * (q % 2) + 8 * (q / (2 * S::rows));
+  }
 
-// The node axis is long, so deriving a channel from a flat element index
-// costs a 64-bit division per thread. The elementwise epilogues instead index
-// (channel, node) directly: threadIdx.x selects the float4 lane within a row,
-// keeping a warp's accesses contiguous, and the remaining dimensions walk the
-// node axis. Layer widths are multiples of four (Python gate).
-struct ElementwiseLaunch {
-  dim3 grid;
-  dim3 block;
+  float* sheet_;
+  int rows_;
+  int rank_;
 };
 
-ElementwiseLaunch elementwise_launch(long n_node, int dout) {
-  const int lanes = dout / 4;
-  const int rows = std::max(1, 256 / lanes);
-  return {dim3((unsigned)ceil_div(n_node, rows)), dim3(lanes, rows)};
+/// C = A @ B over a resident A operand of ``width`` features. With
+/// ``gather``, the A operand is that of an open exchange: the CTA's own
+/// columns are in place, and the slices of the other CTAs stream in from the
+/// exchange sheet in ascending column order, the contraction order, each
+/// loaded while the slices before it are multiplied and stored just before
+/// its first use, so the GEMM waits on the transfer of the first slice only.
+/// Storing the last of them releases the sheet. Without, the operand is
+/// complete already.
+template 
+__device__ __forceinline__ void fused_gathered_gemm(
+    float (*act)[S::act_stride],
+    WeightStream& stream,
+    const PanelExchange& exchange,
+    bool gather,
+    const LaneOrigin& o,
+    Accumulators& acc) {
+  constexpr int kSlices = S::width / S::depth;
+  if constexpr (S::cluster == 1) {
+    fused_resident_gemm(act, stream, kSlices, o, acc);
+  } else {
+    static_assert(S::cols % S::depth == 0,
+                  "every column slice is whole depth slices");
+    constexpr int kPerPeer = S::cols / S::depth;
+    zero_accumulators(acc);
+    typename PanelExchange::Quads incoming;
+    int peer = gather ? exchange.after(-1) : S::cluster;
+    if (peer < S::cluster) {
+      exchange.load(peer, incoming);
+    }
+    for (int s = 0; s < kSlices; ++s) {
+      if (s == peer * kPerPeer) {
+        exchange.store(act, peer, incoming);
+        peer = exchange.after(peer);
+        if (peer < S::cluster) {
+          exchange.load(peer, incoming);
+        } else {
+          exchange.release();
+        }
+      }
+      multiply_slice(act + s * S::depth, stream.next(), o, acc);
+    }
+    __syncthreads();
+  }
 }
 
-// y = act(pre + b) (+ x residual). The pre-activation stays where the GEMM
-// wrote it, so the layer stores one tensor instead of two.
-template 
-__global__ void layer_epilogue_kernel(long n_node,
-                                      int dout,
-                                      const float* __restrict__ pre,
-                                      const float* __restrict__ b,
-                                      const float* __restrict__ x,
-                                      int residual,
-                                      float* __restrict__ y) {
-  const long node = blockIdx.x * (long)blockDim.y + threadIdx.y;
-  if (node >= n_node) {
-    return;
+/// e = (sum of the group partials in ascending order + b_head) +
+/// bias_atom_e[atype] for the ``rows`` rows of the panel at ``row0``.
+template 
+__device__ __forceinline__ void panel_energy(const FusedNetwork& p,
+                                             const float* partial,
+                                             long row0,
+                                             int rows) {
+  for (int r = threadIdx.x; r < rows; r += S::threads) {
+    const long m = row0 + r;
+    p.e[m] = head_energy(partial + r, S::rows, S::width / kHeadGroup, p.b_head,
+                         p.bias_atom_e[p.atype[m]]);
   }
-  const int c = (int)threadIdx.x * 4;
-  const long t = node * dout + c;
-  const float4 p = *reinterpret_cast(pre + t);
-  const float4 bb =
-      b ? *reinterpret_cast(b + c) : make_float4(0, 0, 0, 0);
-  float4 yy =
-      make_float4(act_value(p.x + bb.x), act_value(p.y + bb.y),
-                  act_value(p.z + bb.z), act_value(p.w + bb.w));
-  if (residual) {
-    const float4 xx = *reinterpret_cast(x + t);
-    yy.x += xx.x;
-    yy.y += xx.y;
-    yy.z += xx.z;
-    yy.w += xx.w;
-  }
-  *reinterpret_cast(y + t) = yy;
-}
-
-// Energy head: e[n] = h[n] @ w_head + b_head + bias_atom_e[atype[n]].
-__global__ void head_kernel(long n_node,
-                            int width,
-                            const float* __restrict__ h,
-                            const float* __restrict__ w_head,
-                            const float* __restrict__ b_head,
-                            const double* __restrict__ bias_atom_e,
-                            const long* __restrict__ atype,
-                            double* __restrict__ e) {
-  const long n = blockIdx.x * (long)blockDim.x + threadIdx.x;
-  if (n >= n_node) {
-    return;
+}
+
+template 
+__global__ void __launch_bounds__(S::threads, S::min_blocks)
+    fused_network_kernel(const FusedNetwork p) {
+  extern __shared__ __align__(16) float fused_smem[];
+  auto act = reinterpret_cast(fused_smem);
+  auto ring = reinterpret_cast(
+      fused_smem + S::act_floats);
+  auto xbuf = reinterpret_cast(
+      fused_smem + S::act_floats + S::weight_floats);
+  // (width / 32, rows) partials of the energy head, in the input staging of
+  // the first layer, which the forward has passed by then.
+  float* const partial = &xbuf[0][0][0];
+  const LaneOrigin o = lane_origin();
+  // A cluster is S::cluster consecutive CTAs of the grid, and its CTA
+  // ``rank`` owns the columns from ``col0`` of every GEMM.
+  const int rank = blockIdx.x % S::cluster;
+  const int col0 = rank * S::cols;
+  const long row0 = (long)(blockIdx.x / S::cluster) * S::rows;
+  // Rows of the panel inside the run. Panel rows past it compute on the run's
+  // last row, and in a cluster on zeros in the columns of the other CTAs;
+  // every GEMM keeps rows apart, so their values reach only rows that are
+  // never stored.
+  const int rows = (int)min((long)S::rows, p.nodes - row0);
+  const int last = p.n_layer - 1;
+  const long sheet = p.nodes * S::width;
+  const PanelExchange exchange(p.saved + last * sheet, row0, rows, rank);
+  Accumulators acc;
+
+  // Energy cotangent of accumulator row ``i``; rows past the run carry none.
+  const auto de = [&](int i) {
+    const int r = accumulator_row(o, i);
+    return r < rows ? (float)p.d_e[row0 + r] : 0.f;
+  };
+
+  // === Step 1. Forward through every hidden layer ===
+  WeightStream stream(p, ring, col0);
+  for (int l = 0; l <= last; ++l) {
+    if (l == 0) {
+      FusedInputColumns sa(p.x, row0, p.nodes, p.din);
+      zero_accumulators(acc);
+      const int slices = (p.din + S::depth - 1) / S::depth;
+      sa.fetch(0);
+      sa.commit(xbuf[0]);
+      for (int s = 0; s < slices; ++s) {
+        const bool next = s + 1 < slices;
+        if (next) {
+          sa.fetch((s + 1) * S::depth);
+        }
+        multiply_slice(xbuf[s & 1], stream.next(), o, acc);
+        if (next) {
+          sa.commit(xbuf[(s + 1) & 1]);
+        }
+      }
+      __syncthreads();
+    } else {
+      exchange.open();
+      fused_gathered_gemm(act, stream, exchange, true, o, acc);
+    }
+
+    // Pre-activation, activation and identity residual. The last layer
+    // reduces its activation against the head weight and leaves the seed
+    // cotangent d_e * w_head * act'(pre + b) of the backward in its place.
+    // The parameter quads are loaded before the barrier and the stores of
+    // the epilogue, so that their L2 round trips overlap.
+    float4 bias[S::tn / 4];
+    float4 wh[S::tn / 4];
+#pragma unroll
+    for (int g = 0; g < S::tn / 4; ++g) {
+      const int n = col0 + accumulator_col(o, g);
+      bias[g] = p.b[l] ? load4(p.b[l] + n) : zero4();
+      wh[g] = l == last ? load4(p.w_head + n) : zero4();
+    }
+    exchange.acquire();
+    // Only the residual of the first layer reads the input, and a layer has a
+    // residual only when it is square, so its input rows are whole quads.
+    const bool vector_input = reinterpret_cast(p.x) % 16 == 0;
+    float head[S::tm][S::tn / 4];
+#pragma unroll
+    for (int i = 0; i < S::tm; ++i) {
+      const int r = accumulator_row(o, i);
+      const long mc = row0 + min(r, rows - 1);
+#pragma unroll
+      for (int g = 0; g < S::tn / 4; ++g) {
+        const int n = col0 + accumulator_col(o, g);
+        const float4 z = accumulator_quad(acc, i, g);
+        if (l < last && r < rows) {
+          store4(p.saved + l * sheet + mc * S::width + n, z);
+        }
+        float4 v = act_value4(add4(z, bias[g]));
+        if (p.residual[l]) {
+          float4 rv;
+          if (l > 0) {
+            rv = load_feature_quad(act, n, r);
+          } else if (vector_input) {
+            rv = load4_coherent(p.x + mc * p.din + n);
+          } else {
+            const float* xr = p.x + mc * p.din + n;
+            rv = make_float4(__ldcg(xr), __ldcg(xr + 1), __ldcg(xr + 2),
+                             __ldcg(xr + 3));
+          }
+          v = add4_rn(v, rv);
+        }
+        if (l == last) {
+          float q = v.x * wh[g].x;
+          q = fmaf(v.y, wh[g].y, q);
+          q = fmaf(v.z, wh[g].z, q);
+          head[i][g] = fmaf(v.w, wh[g].w, q);
+          const float4 dh = scale4(de(i), wh[g]);
+          const float4 slope = act_derivative4(add4(z, bias[g]));
+          v = make_float4(dh.x * slope.x, dh.y * slope.y, dh.z * slope.z,
+                          dh.w * slope.w);
+        }
+        store_feature_quad(act, n, r, v);
+        exchange.publish(r, n, v);
+      }
+    }
+    if (l < last) {
+      continue;
+    }
+
+    // === Step 2. Energy head: 32-channel group partials, ascending groups ===
+    // Each CTA reduces the groups of its own columns; the first CTA of a
+    // cluster collects those of the others once the seed exchange opens and
+    // sums every group in order.
+#pragma unroll
+    for (int i = 0; i < S::tm; ++i) {
+      const int r = accumulator_row(o, i);
+#pragma unroll
+      for (int g = 0; g < S::tn / 4; ++g) {
+        float q = head[i][g];
+        q += __shfl_xor_sync(0xffffffffu, q, lane_col_stride(1));
+        q += __shfl_xor_sync(0xffffffffu, q, lane_col_stride(2));
+        q += __shfl_xor_sync(0xffffffffu, q, lane_col_stride(4));
+        if (lane_col(threadIdx.x % 32) == 0) {
+          partial[((col0 + accumulator_col(o, g)) / kHeadGroup) * S::rows +
+                  r] = q;
+        }
+      }
+    }
   }
-  const float* row = h + n * width;
-  float acc = 0.f;
-  for (int k = 0; k < width; k += 4) {
-    const float4 hv = *reinterpret_cast(row + k);
-    const float4 wv = *reinterpret_cast(w_head + k);
-    acc += hv.x * wv.x + hv.y * wv.y + hv.z * wv.z + hv.w * wv.w;
+
+  // === Step 3. Backward from the seed cotangent to the input ===
+  for (int l = last; l >= 0; --l) {
+    const int din = l == 0 ? p.din : S::width;
+    const int passes = (din + S::width - 1) / S::width;
+    for (int pass = 0; pass < passes; ++pass) {
+      const int n0 = pass * S::width + col0;
+      if (pass == 0) {
+        exchange.open();
+        if (l == last) {
+          exchange.collect(partial);
+          __syncthreads();
+          if (rank == 0) {
+            panel_energy(p, partial, row0, rows);
+          }
+        }
+      }
+      fused_gathered_gemm(act, stream, exchange, pass == 0, o, acc);
+
+      // The global operands of the epilogue: the bias of the layer below,
+      // and per quad the factor of the identity residual -- the head weight
+      // at the last layer, which the energy cotangent scales, and below it
+      // the output cotangent that replaced the pre-activation sheet of the
+      // layer -- and the stored pre-activation of the layer below. A tile of
+      // at most 16 accumulators in a CTA of at most 256 threads loads the
+      // operands of every quad before the barrier and its stores, as the
+      // forward loads its parameters, so that their round trips overlap each
+      // other and the wait of a cluster; every other shape loads those of
+      // each quad where it uses them. The two placements compute the same
+      // values; the split keeps every shape within its register budget
+      // without spills on sm_80, sm_90 and sm_120.
+      const float* bb = l > 0 ? p.b[l - 1] : nullptr;
+      float4 bias[S::tn / 4];
+#pragma unroll
+      for (int g = 0; g < S::tn / 4; ++g) {
+        bias[g] = bb ? load4(bb + n0 + accumulator_col(o, g)) : zero4();
+      }
+      float4 factor[S::tm][S::tn / 4];
+      float4 pre[S::tm][S::tn / 4];
+      const auto load_operands = [&](int i, int g) {
+        const int n = n0 + accumulator_col(o, g);
+        const long mc = row0 + min(accumulator_row(o, i), rows - 1);
+        factor[i][g] =
+            !p.residual[l] ? zero4()
+            : l == last
+                ? load4(p.w_head + n)
+                : load4_coherent(p.saved + l * sheet + mc * S::width + n);
+        pre[i][g] =
+            l > 0
+                ? load4_coherent(p.saved + (l - 1) * sheet + mc * S::width + n)
+                : zero4();
+      };
+      constexpr bool kAhead = S::tm * S::tn <= 16 && S::threads <= 256;
+      if constexpr (kAhead) {
+#pragma unroll
+        for (int i = 0; i < S::tm; ++i) {
+#pragma unroll
+          for (int g = 0; g < S::tn / 4; ++g) {
+            load_operands(i, g);
+          }
+        }
+      }
+      if (l > 0) {
+        exchange.acquire();
+      }
+      const bool vector_output =
+          p.din % 4 == 0 && reinterpret_cast(p.d_x) % 16 == 0;
+#pragma unroll
+      for (int i = 0; i < S::tm; ++i) {
+        const int r = accumulator_row(o, i);
+        const long mc = row0 + min(r, rows - 1);
+#pragma unroll
+        for (int g = 0; g < S::tn / 4; ++g) {
+          const int n = n0 + accumulator_col(o, g);
+          if constexpr (!kAhead) {
+            load_operands(i, g);
+          }
+          float4 v = accumulator_quad(acc, i, g);
+          if (p.residual[l]) {
+            // v + d_e * w_head with one rounding, as the per-layer path forms
+            // it; below the last layer the scale of one leaves a plain sum.
+            v = fma4(l == last ? de(i) : 1.f, factor[i][g], v);
+          }
+          if (l > 0) {
+            // dpre of the layer below, from its stored pre-activation, which
+            // the cotangent then replaces.
+            const float4 slope = act_derivative4(add4(pre[i][g], bias[g]));
+            const float4 dpre = make_float4(v.x * slope.x, v.y * slope.y,
+                                            v.z * slope.z, v.w * slope.w);
+            store_feature_quad(act, n, r, dpre);
+            exchange.publish(r, n, dpre);
+            if (r < rows) {
+              store4(p.saved + (l - 1) * sheet + mc * S::width + n, v);
+            }
+          } else if (r < rows && n < p.din) {
+            float* out = p.d_x + mc * p.din + n;
+            if (vector_output) {
+              store4(out, v);
+            } else {
+              const float c[4] = {v.x, v.y, v.z, v.w};
+#pragma unroll
+              for (int q = 0; q < 4; ++q) {
+                if (n + q < p.din) {
+                  out[q] = c[q];
+                }
+              }
+            }
+          }
+        }
+      }
+    }
   }
-  e[n] = (double)(acc + (b_head ? b_head[0] : 0.f)) + bias_atom_e[atype[n]];
+  exchange.acquire();
 }
 
-// Construct the head gradient and apply the final hidden activation VJP in one
-// pass. A square final layer also preserves the unmodified head gradient for
-// its identity branch.
-template 
-__global__ void seed_backward_epilogue_kernel(
-    long n_node,
-    int dout,
-    const double* __restrict__ d_e,
-    const float* __restrict__ w_head,
-    const float* __restrict__ pre,
-    const float* __restrict__ b,
-    float* __restrict__ dpre,
-    float* __restrict__ residual_out) {
-  const long node = blockIdx.x * (long)blockDim.y + threadIdx.y;
-  if (node >= n_node) {
+/// Destinations of the transposed weight sheets of one network, passed by
+/// value so that a captured launch carries them in its parameters.
+struct TransposedSheets {
+  float* wt[kFusedLayers];
+};
+
+/// Edge of the square tiles in which transpose_weights_kernel moves a weight.
+constexpr int kTransposeTile = 32;
+
+/// Transpose every weight of the network into the row-major (width, pitch)
+/// layout the backward of the fused network streams, zero-padded past din.
+/// CTA (x, y, l) of 32 x 8 threads moves the 32 x 32 tile of layer l at
+/// weight rows from 32 * x and columns from 32 * y through shared memory, so
+/// that each warp reads 32 consecutive floats of a weight row and writes 32
+/// consecutive floats of a sheet row; the extra column of the tile keeps the
+/// transposed reads on distinct banks.
+__global__ void transpose_weights_kernel(const FusedNetwork p,
+                                         int width,
+                                         const TransposedSheets sheets) {
+  __shared__ float tile[kTransposeTile][kTransposeTile + 1];
+  const int l = blockIdx.z;
+  const int din = l == 0 ? p.din : width;
+  const int pitch = p.pitch[l];
+  const int n0 = blockIdx.x * kTransposeTile;
+  const int k0 = blockIdx.y * kTransposeTile;
+  if (n0 >= pitch) {
     return;
   }
-  const int c = (int)threadIdx.x * 4;
-  const long t = node * dout + c;
-  const float de = (float)d_e[node];
-  const float4 wv = *reinterpret_cast(w_head + c);
-  const float4 dh = make_float4(de * wv.x, de * wv.y, de * wv.z, de * wv.w);
-  const float4 p = *reinterpret_cast(pre + t);
-  const float4 bb =
-      b ? *reinterpret_cast(b + c) : make_float4(0, 0, 0, 0);
-  if (residual_out) {
-    *reinterpret_cast(residual_out + t) = dh;
-  }
-  *reinterpret_cast(dpre + t) =
-      make_float4(dh.x * act_derivative(p.x + bb.x),
-                  dh.y * act_derivative(p.y + bb.y),
-                  dh.z * act_derivative(p.z + bb.z),
-                  dh.w * act_derivative(p.w + bb.w));
-}
-
-// Convert dh to dpre in place: dh *= act'(pre + b). A square residual layer
-// also preserves the unmodified dh in its output buffer before the GEMM
-// accumulates the weighted branch with beta = 1.
-template 
-__global__ void backward_epilogue_kernel(long n_node,
-                                         int dout,
-                                         const float* __restrict__ dh,
-                                         const float* __restrict__ pre,
-                                         const float* __restrict__ b,
-                                         float* __restrict__ dpre,
-                                         float* __restrict__ residual_out) {
-  const long node = blockIdx.x * (long)blockDim.y + threadIdx.y;
-  if (node >= n_node) {
-    return;
+  for (int i = threadIdx.y; i < kTransposeTile; i += blockDim.y) {
+    const int n = n0 + i;
+    const int k = k0 + threadIdx.x;
+    tile[i][threadIdx.x] =
+        n < din && k < width ? p.w[l][(long)n * width + k] : 0.f;
   }
-  const int c = (int)threadIdx.x * 4;
-  const long t = node * dout + c;
-  const float4 d = *reinterpret_cast(dh + t);
-  const float4 p = *reinterpret_cast(pre + t);
-  const float4 bb =
-      b ? *reinterpret_cast(b + c) : make_float4(0, 0, 0, 0);
-  if (residual_out) {
-    *reinterpret_cast(residual_out + t) = d;
+  __syncthreads();
+  for (int i = threadIdx.y; i < kTransposeTile; i += blockDim.y) {
+    const int k = k0 + i;
+    const int n = n0 + threadIdx.x;
+    if (k < width && n < pitch) {
+      sheets.wt[l][(long)k * pitch + n] = tile[threadIdx.x][i];
+    }
   }
-  *reinterpret_cast(dpre + t) =
-      make_float4(d.x * act_derivative(p.x + bb.x),
-                  d.y * act_derivative(p.y + bb.y),
-                  d.z * act_derivative(p.z + bb.z),
-                  d.w * act_derivative(p.w + bb.w));
 }
 
+// ---------------------------------------------------------------------------
+// Launch.
+// ---------------------------------------------------------------------------
+
 /// Dispatch a kernel template over the two supported activation codes.
 template 
 void dispatch_activation(long act, Fn&& launch) {
@@ -274,6 +1889,94 @@ void dispatch_activation(long act, Fn&& launch) {
   }
 }
 
+/// Dispatch over the two tiles. A run small enough to leave the large tile's
+/// grid short of a few waves takes the small tile; the threshold of four and
+/// a half small CTAs per SM is where the two break even on the fitting shapes
+/// of the DPA4C presets on sm_120.
+template 
+void dispatch_tile(long nodes, int width, Fn&& launch) {
+  const long sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  const long small_ctas =
+      ceil_div(nodes, SmallTile::rows) * ceil_div(width, SmallTile::cols);
+  if (2 * small_ctas <= 9 * sms) {
+    launch(SmallTile{});
+  } else {
+    launch(LargeTile{});
+  }
+}
+
+template 
+unsigned grid_size(long nodes, int width) {
+  return (unsigned)(ceil_div(nodes, Tile::rows) * ceil_div(width, Tile::cols));
+}
+
+/// Column tiles of a GEMM with ``width`` outputs under its routed tile.
+long column_tiles(long nodes, int width) {
+  long tiles = 0;
+  dispatch_tile(nodes, width, [&](auto tile) {
+    tiles = ceil_div(width, decltype(tile)::cols);
+  });
+  return tiles;
+}
+
+/// Launch one forward layer. Returns whether the energy still needs
+/// head_finalize_kernel, which is the case when the last layer is wider than
+/// a tile.
+bool launch_layer_forward(cudaStream_t stream,
+                          const LayerForward& p,
+                          long act) {
+  bool finalize = false;
+  dispatch_tile(p.nodes, p.dout, [&](auto tile) {
+    using Tile = decltype(tile);
+    const unsigned grid = grid_size(p.nodes, p.dout);
+    const bool vector =
+        p.din % 4 == 0 && reinterpret_cast(p.x) % 16 == 0;
+    dispatch_activation(act, [&](auto tag) {
+      constexpr int kAct = decltype(tag)::value;
+      if (vector) {
+        layer_forward_kernel
+            <<>>(p);
+      } else {
+        layer_forward_kernel
+            <<>>(p);
+      }
+    });
+    finalize = p.y == nullptr && p.dout > Tile::cols;
+  });
+  return finalize;
+}
+
+void launch_layer_backward(cudaStream_t stream,
+                           const LayerBackward& p,
+                           Cotangent source,
+                           long act) {
+  // An input of zero width has no cotangent to compute.
+  if (p.din == 0) {
+    return;
+  }
+  dispatch_tile(p.nodes, p.din, [&](auto tile) {
+    using Tile = decltype(tile);
+    const unsigned grid = grid_size(p.nodes, p.din);
+    dispatch_activation(act, [&](auto tag) {
+      constexpr int kAct = decltype(tag)::value;
+      switch (source) {
+        case Cotangent::kHead:
+          layer_backward_kernel
+              <<>>(p);
+          break;
+        case Cotangent::kChain:
+          layer_backward_kernel
+              <<>>(p);
+          break;
+        case Cotangent::kFormed:
+          layer_backward_kernel
+              <<>>(p);
+          break;
+      }
+    });
+  });
+}
+
 }  // namespace
 
 namespace {
@@ -307,12 +2010,96 @@ FittingLayerPlan validate_fitting_forward_inputs(
   return plan;
 }
 
+/// Whether `t` is a contiguous tensor of `dtype` on the current CUDA device,
+/// or empty when `optional`.
+bool resident_parameter(const torch::Tensor& t,
+                        at::ScalarType dtype,
+                        bool optional) {
+  if (optional && t.numel() == 0) {
+    return true;
+  }
+  return t.is_cuda() && t.get_device() == c10::cuda::current_device() &&
+         t.is_contiguous() && t.scalar_type() == dtype;
+}
+
+bool aligned_parameter(const torch::Tensor& t) {
+  return t.numel() == 0 ||
+         reinterpret_cast(t.data_ptr()) % 16 == 0;
+}
+
+/// The kernels read the weight of every layer as aligned float4 from a
+/// contiguous fp32 (din, dout) matrix on the device they run on, the layers
+/// chain, and every hidden width is a positive multiple of four.
+void check_weights(const std::vector& ws) {
+  TORCH_CHECK(!ws.empty(), "graph_fitting: the network has no layer");
+  for (size_t l = 0; l < ws.size(); ++l) {
+    TORCH_CHECK(resident_parameter(ws[l], torch::kFloat32, false) &&
+                    aligned_parameter(ws[l]) && ws[l].dim() == 2,
+                "graph_fitting: the weight of layer ", l,
+                " must be a contiguous, 16-byte aligned fp32 matrix on the "
+                "current CUDA device");
+    TORCH_CHECK(ws[l].size(1) > 0 && ws[l].size(1) % 4 == 0,
+                "graph_fitting: the width of layer ", l,
+                " must be a positive multiple of four, got ", ws[l].size(1));
+    TORCH_CHECK(l == 0 || ws[l].size(0) == ws[l - 1].size(1),
+                "graph_fitting: layer ", l, " takes ", ws[l].size(0),
+                " features, but layer ", l - 1, " yields ", ws[l - 1].size(1));
+  }
+}
+
+/// `check_weights` together with the biases, the residual flags and the head
+/// weight, which the kernels read like the weights.
+void check_parameters(const std::vector& ws,
+                      const std::vector& bs,
+                      const std::vector& resnets,
+                      const torch::Tensor& w_head) {
+  check_weights(ws);
+  TORCH_CHECK(bs.size() == ws.size() && resnets.size() == ws.size(),
+              "graph_fitting: expected one bias and one residual flag per "
+              "fitting layer, got ",
+              bs.size(), " and ", resnets.size(), " for ", ws.size(),
+              " layers");
+  for (size_t l = 0; l < ws.size(); ++l) {
+    TORCH_CHECK(resident_parameter(bs[l], torch::kFloat32, true) &&
+                    aligned_parameter(bs[l]) &&
+                    (bs[l].numel() == 0 || bs[l].numel() == ws[l].size(1)),
+                "graph_fitting: the bias of layer ", l,
+                " must be empty or a contiguous, 16-byte aligned fp32 vector "
+                "of the layer width on the current CUDA device");
+  }
+  TORCH_CHECK(resident_parameter(w_head, torch::kFloat32, false) &&
+                  aligned_parameter(w_head) &&
+                  w_head.numel() == ws.back().size(1),
+              "graph_fitting: the head weight must be a contiguous, 16-byte "
+              "aligned fp32 tensor of ",
+              ws.back().size(1), " entries on the current CUDA device");
+}
+
+/// `check_parameters` together with the head bias and the per-type energy
+/// bias of a pass that forms the energy.
+void check_energy_parameters(const std::vector& ws,
+                             const std::vector& bs,
+                             const std::vector& resnets,
+                             const torch::Tensor& w_head,
+                             const torch::Tensor& b_head,
+                             const torch::Tensor& bias_atom_e) {
+  check_parameters(ws, bs, resnets, w_head);
+  TORCH_CHECK(
+      resident_parameter(b_head, torch::kFloat32, true) && b_head.numel() <= 1,
+      "graph_fitting: the head bias must be empty or a single fp32 "
+      "value on the current CUDA device");
+  TORCH_CHECK(bias_atom_e.dim() == 1 &&
+                  resident_parameter(bias_atom_e, torch::kFloat64, false),
+              "graph_fitting: bias_atom_e must be a contiguous fp64 vector on "
+              "the current CUDA device");
+}
+
 }  // namespace
 
 // Evaluate the network over one contiguous run of nodes. Every full-width
 // tensor is indexed from the run's first node, so the same code serves the
 // whole node axis and a single tile of it. ``saved`` and ``activation`` are
-// sized for the run, not for the system.
+// 16-byte aligned and sized for the run, not for the system.
 void fitting_forward_range(cudaStream_t stream,
                            const FittingLayerPlan& plan,
                            const float* x,
@@ -329,37 +2116,49 @@ void fitting_forward_range(cudaStream_t stream,
                            float* saved,
                            float* const activation[2],
                            double* e) {
+  check_energy_parameters(ws, bs, resnets, w_head, b_head, bias_atom_e);
   const float* cur = x;
   int din = (int)input_width;
+  LayerForward p{};
+  p.w_head = w_head.data_ptr();
+  p.b_head = b_head.numel() ? b_head.data_ptr() : nullptr;
+  p.bias_atom_e = bias_atom_e.data_ptr();
+  p.atype = atype;
+  p.e = e;
+  p.nodes = run_nodes;
+  bool finalize = false;
   for (int l = 0; l < plan.n_layer; ++l) {
     const int dout = (int)ws[l].size(1);
-    float* pre = saved + plan.offset[l] * run_nodes;
-    float* y = activation[l & 1];
-    gemm_nn(stream, cur, ws[l].data_ptr(), pre, (int)run_nodes, dout,
-            din);
-    const ElementwiseLaunch shape = elementwise_launch(run_nodes, dout);
-    const bool residual = resnets[l] && dout == din;
-    dispatch_activation(act, [&](auto tag) {
-      layer_epilogue_kernel
-          <<>>(
-              run_nodes, dout, pre,
-              bs[l].numel() ? bs[l].data_ptr() : nullptr, cur,
-              residual ? 1 : 0, y);
-    });
+    const bool last = l == plan.n_layer - 1;
+    p.x = cur;
+    p.w = ws[l].data_ptr();
+    p.b = bs[l].numel() ? bs[l].data_ptr() : nullptr;
+    p.pre = saved + plan.offset[l] * run_nodes;
+    // The last layer stores no activation; its slot, which it does not read,
+    // holds the partial head dots instead.
+    p.y = last ? nullptr : activation[l & 1];
+    p.head_partial = last ? activation[l & 1] : nullptr;
+    p.din = din;
+    p.dout = dout;
+    p.residual = resnets[l] && dout == din ? 1 : 0;
+    finalize = launch_layer_forward(stream, p, act);
     FITTING_CHECK_LAUNCH("graph_fitting layer");
-    cur = y;
+    cur = p.y;
     din = dout;
   }
-  head_kernel<<>>(
-      run_nodes, din, cur, w_head.data_ptr(),
-      b_head.numel() ? b_head.data_ptr() : nullptr,
-      bias_atom_e.data_ptr(), atype, e);
-  FITTING_CHECK_LAUNCH("graph_fitting head");
+  if (finalize) {
+    head_finalize_kernel<<>>(
+        run_nodes, (int)ceil_div(din, kHeadGroup), p.head_partial, p.b_head,
+        p.bias_atom_e, atype, e);
+    FITTING_CHECK_LAUNCH("graph_fitting head");
+  }
 }
 
 // Propagate the head cotangent of one run of nodes back to the input. ``dh``
-// and ``dh_next`` are scratch of the run's size; ``d_x`` is indexed from the
-// run's first node.
+// and ``dh_next`` are 16-byte aligned (run, width_max) scratch, written from
+// the last layer down: ``dh`` is needed by networks of two or more layers and
+// ``dh_next`` by networks of three or more. ``d_x`` is indexed from the run's
+// first node.
 void fitting_backward_range(cudaStream_t stream,
                             const FittingLayerPlan& plan,
                             const double* d_e,
@@ -373,33 +2172,789 @@ void fitting_backward_range(cudaStream_t stream,
                             float* dh,
                             float* dh_next,
                             float* d_x) {
+  check_parameters(ws, bs, resnets, w_head);
+  // Layer l writes its input cotangent to the slot its successor did not
+  // write, or to d_x for the first layer. Above a layer without identity
+  // residual whose GEMM spans several column tiles it writes that layer's dpre
+  // instead, so the dpre is formed once rather than once per column tile. A
+  // GEMM of a single column tile forms its dpre once in its pipelined main
+  // loop, which costs less than forming it in this epilogue.
+  const auto residual = [&](int l) {
+    return resnets[l] && ws[l].size(0) == ws[l].size(1);
+  };
+  float* const slot[2] = {dh, dh_next};
+  LayerBackward p{};
+  p.d_e = d_e;
+  p.w_head = w_head.data_ptr();
+  p.nodes = run_nodes;
+  Cotangent source = Cotangent::kHead;
   for (int l = plan.n_layer - 1; l >= 0; --l) {
-    const int dout = (int)ws[l].size(1);
-    const int din = (int)ws[l].size(0);
-    const float* pre = saved + plan.offset[l] * run_nodes;
-    const float* b = bs[l].numel() ? bs[l].data_ptr() : nullptr;
-    float* out = l > 0 ? dh_next : d_x;
-    const bool residual = resnets[l] && dout == din;
-    const ElementwiseLaunch shape = elementwise_launch(run_nodes, dout);
+    const bool emit = l > 0 && !residual(l - 1) &&
+                      column_tiles(run_nodes, (int)ws[l - 1].size(0)) > 1;
+    p.pre = saved + plan.offset[l] * run_nodes;
+    p.b = bs[l].numel() ? bs[l].data_ptr() : nullptr;
+    p.w = ws[l].data_ptr();
+    p.out = l > 0 ? slot[(plan.n_layer - 1 - l) & 1] : d_x;
+    p.below_pre = emit ? saved + plan.offset[l - 1] * run_nodes : nullptr;
+    p.below_b =
+        emit && bs[l - 1].numel() ? bs[l - 1].data_ptr() : nullptr;
+    p.din = (int)ws[l].size(0);
+    p.dout = (int)ws[l].size(1);
+    p.residual = residual(l) ? 1 : 0;
+    launch_layer_backward(stream, p, source, act);
+    FITTING_CHECK_LAUNCH("graph_fitting backward layer");
+    p.dh = p.out;
+    source = emit ? Cotangent::kFormed : Cotangent::kChain;
+  }
+}
+
+namespace {
+
+/// Launch configuration of ``panels`` panels under shape S: one CTA per
+/// panel, in clusters of S::cluster consecutive CTAs. ``cluster`` holds the
+/// cluster attribute that the configuration points to.
+template 
+cudaLaunchConfig_t fused_launch_config(long panels,
+                                       cudaStream_t stream,
+                                       cudaLaunchAttribute* cluster) {
+  cluster->id = cudaLaunchAttributeClusterDimension;
+  cluster->val.clusterDim.x = S::cluster;
+  cluster->val.clusterDim.y = 1;
+  cluster->val.clusterDim.z = 1;
+  cudaLaunchConfig_t config{};
+  config.gridDim = dim3((unsigned)(panels * S::cluster));
+  config.blockDim = dim3(S::threads);
+  config.dynamicSmemBytes = S::bytes;
+  config.stream = stream;
+  config.attrs = cluster;
+  config.numAttrs = S::cluster > 1 ? 1 : 0;
+  return config;
+}
+
+/// Opt the kernel of shape S into its dynamic shared memory and, for a
+/// cluster above the portable size of eight, into non-portable cluster sizes,
+/// up to the largest cluster that the device launches.
+template 
+cudaError_t set_fused_attributes() {
+  const auto kernel = fused_network_kernel;
+  if constexpr (S::cluster > 8) {
+    const cudaError_t error = cudaFuncSetAttribute(
+        kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1);
+    if (error != cudaSuccess) {
+      return error;
+    }
+  }
+  return cudaFuncSetAttribute(
+      kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)S::bytes);
+}
+
+/// `set_fused_attributes` for a launch, which fails when the device refuses
+/// them.
+template 
+void configure_fused_kernel() {
+  const cudaError_t error = set_fused_attributes();
+  if (error != cudaSuccess) {
+    // The refusal is cleared, so that no later launch check reports it.
+    (void)cudaGetLastError();
+    TORCH_CHECK(false,
+                "graph_fitting fused network: cannot configure "
+                "clusters of ",
+                S::cluster, " CTAs with ", S::bytes,
+                " bytes of dynamic shared memory: ", cudaGetErrorString(error));
+  }
+}
+
+/// Whether a query of the fused kernel succeeded. A refused query leaves its
+/// error behind, which is cleared, so that the shape it asked about is merely
+/// not used.
+bool accepted(cudaError_t error) {
+  if (error == cudaSuccess) {
+    return true;
+  }
+  (void)cudaGetLastError();
+  return false;
+}
+
+/// What the current device runs of shape S, which no run changes and which
+/// each device therefore answers once.
+struct FusedCapacity {
+  /// Lower virtual architecture of the code that the device runs for the two
+  /// activations: that of the native code it selected, which may be an
+  /// earlier minor revision of its own (compute_80 on sm_86), or the PTX
+  /// target that the driver compiled from.
+  int version;
+  /// Clusters of the shape that the device keeps resident at once for both
+  /// activations; zero for a single-CTA shape, and zero unless the device
+  /// launches clusters, the images were compiled for compute_90 or later, and
+  /// it launches clusters of S::cluster CTAs of this kernel, which a device or
+  /// partition with fewer multiprocessors per GPC may not. The runtime counts
+  /// the resident clusters against the multiprocessors of each GPC and seats
+  /// several CTAs on one multiprocessor where their registers and shared
+  /// memory allow, which the narrow network does.
+  int clusters;
+  /// CTAs of a single-CTA shape that a multiprocessor keeps resident at once
+  /// for both activations, which sets the waves in which a run of its panels
+  /// passes; zero for a cluster shape and where the device refuses the query.
+  int resident;
+};
+
+template 
+FusedCapacity fused_capacity() {
+  static std::mutex mutex;
+  static std::map capacities;
+  const c10::DeviceIndex device = c10::cuda::current_device();
+  const std::lock_guard lock(mutex);
+  const auto known = capacities.find(device);
+  if (known != capacities.end()) {
+    return known->second;
+  }
+  // A query the device or the driver refuses reads as a version of zero or as
+  // no resident cluster, so the shape is not used and the per-layer kernels
+  // serve the run instead of the call failing.
+  FusedCapacity capacity{std::numeric_limits::max(), 0,
+                         S::cluster == 1 ? std::numeric_limits::max() : 0};
+  for (const long act : {0L, 1L}) {
     dispatch_activation(act, [&](auto tag) {
       constexpr int kAct = decltype(tag)::value;
-      if (l == plan.n_layer - 1) {
-        seed_backward_epilogue_kernel
-            <<>>(
-                run_nodes, dout, d_e, w_head.data_ptr(), pre, b, dh,
-                residual ? out : nullptr);
-      } else {
-        backward_epilogue_kernel<<>>(
-            run_nodes, dout, dh, pre, b, dh, residual ? out : nullptr);
+      const auto kernel = fused_network_kernel;
+      cudaFuncAttributes attributes{};
+      const bool known = accepted(cudaFuncGetAttributes(&attributes, kernel));
+      capacity.version =
+          std::min(capacity.version, known ? attributes.ptxVersion : 0);
+      if constexpr (S::cluster == 1) {
+        int resident = 0;
+        if (!accepted(set_fused_attributes()) ||
+            !accepted(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+                &resident, kernel, S::threads, S::bytes))) {
+          resident = 0;
+        }
+        capacity.resident = std::min(capacity.resident, resident);
       }
     });
-    FITTING_CHECK_LAUNCH("graph_fitting backward layer");
-    gemm_nt(stream, dh, ws[l].data_ptr(), out, (int)run_nodes, din, dout,
-            residual ? 1.f : 0.f);
-    if (l > 0) {
-      std::swap(dh, dh_next);
+  }
+  int cluster_launch = 0;
+  const bool launches_clusters =
+      accepted(cudaDeviceGetAttribute(&cluster_launch, cudaDevAttrClusterLaunch,
+                                      device)) &&
+      cluster_launch != 0;
+  if (S::cluster > 1 && launches_clusters && capacity.version >= 90) {
+    capacity.clusters = std::numeric_limits::max();
+    for (const long act : {0L, 1L}) {
+      dispatch_activation(act, [&](auto tag) {
+        constexpr int kAct = decltype(tag)::value;
+        const auto kernel = fused_network_kernel;
+        cudaLaunchAttribute cluster;
+        const cudaLaunchConfig_t config =
+            fused_launch_config(1, nullptr, &cluster);
+        int largest = 0;
+        int clusters = 0;
+        if (!accepted(set_fused_attributes()) ||
+            !accepted(cudaOccupancyMaxPotentialClusterSize(&largest, kernel,
+                                                           &config)) ||
+            largest < S::cluster ||
+            !accepted(
+                cudaOccupancyMaxActiveClusters(&clusters, kernel, &config))) {
+          clusters = 0;
+        }
+        capacity.clusters = std::min(capacity.clusters, clusters);
+      });
+    }
+  }
+  capacities.emplace(device, capacity);
+  return capacity;
+}
+
+/// Launch the first of ``shapes`` that takes a run of ``panels`` panels: of
+/// the shapes whose dynamic shared memory the device lets one CTA opt in to,
+/// a cluster shape when the device keeps a cluster per panel resident at once,
+/// a single-CTA shape unconditionally. Returns whether a shape launched.
+template 
+bool launch_first_fit(long panels, Fn&& launch, Shapes... shapes) {
+  const auto take = [&](auto shape) {
+    using S = decltype(shape);
+    if (S::bytes >
+        at::cuda::getCurrentDeviceProperties()->sharedMemPerBlockOptin) {
+      return false;
+    }
+    if constexpr (S::cluster > 1) {
+      if (fused_capacity().clusters < panels) {
+        return false;
+      }
+    }
+    launch(shape);
+    return true;
+  };
+  return (take(shapes) || ...);
+}
+
+/// A list of fused shapes.
+template 
+struct ShapeList {};
+
+/// launch_first_fit over the shapes of a list.
+template 
+bool launch_first_listed(long panels, Fn&& launch, ShapeList) {
+  return launch_first_fit(panels, launch, Shapes{}...);
+}
+
+/// The fused shapes compiled for a hidden width, in the order of the dispatch:
+/// the latency shapes, the largest cluster first, then the throughput shapes.
+template 
+struct FusedShapes;
+template <>
+struct FusedShapes<64> {
+  using Latency = ShapeList,
+                            FusedShape<64, 1, 4, 4, 2, 2, 1, 8>>;
+  using Throughput = ShapeList,
+                               FusedShape<64, 1, 4, 4, 4, 2, 2, 2, 16>>;
+};
+template <>
+struct FusedShapes<96> {
+  using Latency = ShapeList,
+                            FusedShape<96, 1, 4, 4, 2, 3, 1, 8>>;
+  using WidePanels = FusedShape<96, 1, 4, 4, 4, 3, 2, 2, 16>;
+  using NarrowPanels = FusedShape<96, 1, 4, 4, 2, 3, 3, 3>;
+  using Throughput = ShapeList;
+};
+template <>
+struct FusedShapes<192> {
+  using Latency = ShapeList,
+                            FusedShape<192, 3, 4, 4, 2, 2, 1, 4, 32>,
+                            FusedShape<192, 2, 4, 4, 2, 3, 1, 4, 16>,
+                            FusedShape<192, 1, 4, 4, 2, 6, 1, 8>>;
+  using Throughput = ShapeList>;
+};
+template <>
+struct FusedShapes<256> {
+  using Latency = ShapeList,
+                            FusedShape<256, 4, 4, 4, 2, 2, 1, 4, 32>,
+                            FusedShape<256, 2, 4, 4, 2, 4, 1, 4, 16>,
+                            FusedShape<256, 1, 4, 4, 2, 8, 1, 7, 8, 3>>;
+  using Throughput = ShapeList>;
+};
+template <>
+struct FusedShapes<384> {
+  using Latency = ShapeList,
+                            FusedShape<384, 6, 4, 4, 2, 2, 1, 4, 32>,
+                            FusedShape<384, 4, 4, 4, 2, 3, 1, 4, 16>,
+                            FusedShape<384, 3, 4, 4, 2, 4, 1, 4, 16>,
+                            FusedShape<384, 2, 4, 4, 2, 6, 1, 3, 16>,
+                            FusedShape<384, 1, 8, 4, 1, 12, 1, 3, 8>>;
+  using Throughput = ShapeList<>;
+};
+
+/// Call ``fn`` with the FusedShapes of ``width``; false for a width that no
+/// fused shape is compiled for.
+template 
+bool with_fused_shapes(long width, Fn&& fn) {
+  switch (width) {
+    case 64:
+      return fn(FusedShapes<64>{});
+    case 96:
+      return fn(FusedShapes<96>{});
+    case 192:
+      return fn(FusedShapes<192>{});
+    case 256:
+      return fn(FusedShapes<256>{});
+    case 384:
+      return fn(FusedShapes<384>{});
+    default:
+      return false;
+  }
+}
+
+/// Call ``visit`` with the shapes of a list in order until it returns true;
+/// returns whether it did.
+template 
+bool visit_shapes(ShapeList, [[maybe_unused]] Fn&& visit) {
+  return (visit(Shapes{}) || ...);
+}
+
+/// Call ``visit`` with every fused shape compiled for ``width`` until it
+/// returns true; returns whether it did.
+template 
+bool visit_fused_shapes(long width, Fn&& visit) {
+  return with_fused_shapes(width, [&](auto shapes) {
+    using Shapes = decltype(shapes);
+    return visit_shapes(typename Shapes::Latency{}, visit) ||
+           visit_shapes(typename Shapes::Throughput{}, visit);
+  });
+}
+
+/// Call ``visit`` with the fused shapes of ``width`` that a measurement of
+/// the fitting routes compares on runs of ``latency`` or of more panels than
+/// multiprocessors: the latency shapes on the former, and on the latter the
+/// throughput shapes, or the single-CTA latency shapes of a width that has
+/// none; until it returns true. Returns whether it did.
+template 
+bool visit_compared_shapes(long width, bool latency, Fn&& visit) {
+  return with_fused_shapes(width, [&](auto shapes) {
+    using Shapes = decltype(shapes);
+    if (latency) {
+      return visit_shapes(typename Shapes::Latency{}, visit);
+    }
+    if constexpr (std::is_same_v>) {
+      return visit_shapes(typename Shapes::Latency{}, [&](auto shape) {
+        return decltype(shape)::cluster == 1 && visit(shape);
+      });
+    } else {
+      return visit_shapes(typename Shapes::Throughput{}, visit);
+    }
+  });
+}
+
+/// Dispatch the fused network over the hidden widths it is compiled for.
+/// Returns whether it passed a shape to ``launch``.
+///
+/// A run of more panels than multiprocessors takes the throughput shape of its
+/// width, the fastest on sm_120 of those that fit the register and
+/// shared-memory budgets: the narrow widths take 32-node panels, four CTAs
+/// per SM at a width of 64, three at 96 and two at 192, whose small lane
+/// tiles keep the most panels in flight, and a width of 256 takes 64-node
+/// panels on one CTA per SM, which stream each weight slice to half as many
+/// CTAs as 32-node panels and so keep the L2 traffic within the FMA rate. A
+/// width of 96 moves on to 64-node panels of twelve warps, two CTAs per SM in
+/// place of three, once a run fills two waves of them, or once they pass it in
+/// fewer waves than the 32-node panels, each shape at the CTAs per SM that the
+/// device keeps resident. Such a panel streams each weight slice once for twice
+/// the nodes of a 32-node panel, and on sm_120 the larger panels take 6% less
+/// time at 100,800 nodes and 11% less at 131,072, the default node tile; from
+/// two waves on, this gain outweighs the cost of their last, partly filled
+/// wave, and below two waves a wave fewer outweighs it: on the 92
+/// multiprocessors of an L20, one wave of 64-node panels takes 0.064 ms at
+/// 10,296 nodes against 0.078 ms for two waves of 32-node panels.
+///
+/// A width that the 128 columns of the per-layer large tile divide is left to
+/// the per-layer kernels at runs of more panels than multiprocessors on a
+/// device whose fp32 ridge lies below width / 4, the FLOP per byte of a hidden
+/// layer of those kernels: bound there by the arithmetic rather than by the
+/// activation traffic that the fusion saves, they use more of the fp32 peak
+/// than the fused panels, and at a width of 256 on an H20, whose ridge is 10,
+/// they shorten the LAMMPS step by 1% to 6%. At a width of 384 the per-layer
+/// kernels, whose 64 x 128 tiles reuse each weight slice across twice the
+/// nodes, serve such a run on every device: a 32-node panel of that width
+/// takes 94,848 bytes of shared memory, a multiprocessor of sm_120 to itself,
+/// so the run would pass in waves of one panel per multiprocessor.
+///
+/// A smaller run splits its 32-node panels over the largest cluster, of the
+/// divisors of width / 32 above one, of which the device keeps one per panel
+/// resident at once; the cluster divides the FMAs of each panel between its
+/// CTAs. A cluster above the portable size of eight takes the non-portable
+/// cluster attribute and runs where the device launches clusters of its size.
+/// A run too large for any cluster in one wave, or on a device without
+/// clusters, takes the latency shape: 32-node panels on the most lanes their
+/// tile allows, each alone on its multiprocessor, with a ring deep enough to
+/// keep the weight stream ahead of the FMAs without other panels to hide its
+/// latency. The ring runs ahead through asynchronous copies, so
+/// fitting_fused_serves leaves such a run to the per-layer kernels on a kernel
+/// image of compute capability 7.x, whose copies are plain loads and stores.
+/// The tail of a tiled call that graph_fitting_energy_gradient serves with the
+/// fused network takes such a shape nevertheless, since that operator decides
+/// once for the runs of the call; the fused energy-force operator of DPA4C
+/// decides the tail by its own size (fitting_route).
+template 
+bool dispatch_fused_shape(long width, long run_nodes, Fn&& launch) {
+  const long sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  const long panels = ceil_div(run_nodes, 32);
+  const bool latency = panels <= sms;
+  if (!latency && width % LargeTile::cols == 0 && fp32_ridge() < width / 4.0) {
+    return false;
+  }
+  return with_fused_shapes(width, [&](auto shapes) {
+    using Shapes = decltype(shapes);
+    if (latency) {
+      return launch_first_listed(panels, launch, typename Shapes::Latency{});
+    }
+    if constexpr (std::is_same_v>) {
+      using WidePanels = typename Shapes::WidePanels;
+      using NarrowPanels = typename Shapes::NarrowPanels;
+      // Panels of each shape that one wave of the device holds.
+      const long wide_wave = fused_capacity().resident * sms;
+      const long narrow_wave = fused_capacity().resident * sms;
+      if (wide_wave > 0 &&
+          (run_nodes >= 2 * wide_wave * WidePanels::rows ||
+           (narrow_wave > 0 &&
+            ceil_div(ceil_div(run_nodes, WidePanels::rows), wide_wave) <
+                ceil_div(ceil_div(run_nodes, NarrowPanels::rows),
+                         narrow_wave)))) {
+        return launch_first_fit(panels, launch, WidePanels{});
+      }
+      return launch_first_fit(panels, launch, NarrowPanels{});
+    } else {
+      return launch_first_listed(panels, launch, typename Shapes::Throughput{});
+    }
+  });
+}
+
+/// Whether the device runs shape S on a run of ``run_nodes`` nodes: an image
+/// of compute_70 or later whose shared memory the device admits and, for a
+/// cluster shape, a cluster per panel resident at once. A single-CTA shape
+/// runs a run of any size, in as many waves as it takes.
+template 
+bool fused_shape_runs(long run_nodes) {
+  const FusedCapacity capacity = fused_capacity();
+  if (capacity.version < 70) {
+    return false;
+  }
+  if constexpr (S::cluster > 1) {
+    return capacity.clusters >= ceil_div(run_nodes, S::rows);
+  } else {
+    return capacity.resident > 0;
+  }
+}
+
+/// CTAs of the per-layer kernels of ``Tile`` that a multiprocessor of the
+/// current device keeps resident, the fewest over the forward and the backward
+/// kernels and both activations. Each device answers once.
+template 
+int tile_resident() {
+  static std::mutex mutex;
+  static std::map residents;
+  const c10::DeviceIndex device = c10::cuda::current_device();
+  const std::lock_guard lock(mutex);
+  const auto known = residents.find(device);
+  if (known != residents.end()) {
+    return known->second;
+  }
+  int resident = std::numeric_limits::max();
+  for (const long act : {0L, 1L}) {
+    dispatch_activation(act, [&](auto tag) {
+      constexpr int kAct = decltype(tag)::value;
+      int forward = 0;
+      int backward = 0;
+      if (!accepted(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+              &forward, layer_forward_kernel, Tile::threads,
+              0))) {
+        forward = 0;
+      }
+      if (!accepted(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+              &backward, layer_backward_kernel,
+              Tile::threads, 0))) {
+        backward = 0;
+      }
+      resident = std::min({resident, forward, backward});
+    });
+  }
+  residents.emplace(device, resident);
+  return resident;
+}
+
+// A fitting-route table records which route serves each run size of one
+// network on one device, as measured at freeze time; the fused energy-force
+// operator of DPA4C receives it as a literal argument of the frozen program.
+// It holds kRouteTableMagic, the fingerprint of the device it was measured on
+// (device_fingerprint), the number n of its intervals of run sizes, the
+// inclusive upper bounds of the first n - 1 intervals in increasing order, and
+// one choice per interval: kFittingPerLayer or the key of a fused shape
+// (FusedShape::key) where the measurement departs from the rules, and
+// kFittingRules, which leaves the run to them, everywhere else.
+
+/// First word of a fitting-route table, which also names its layout.
+constexpr int64_t kRouteTableMagic = 0x52545431;
+/// Words of the device fingerprint of a fitting-route table.
+constexpr size_t kFingerprintWords = 11;
+
+/// What sets the waves and the speed of the fitting kernels on the current
+/// device: its compute capability, multiprocessors, L2 size, registers and
+/// shared memory per multiprocessor, shared memory a block may opt in to,
+/// thread and block limits per multiprocessor, memory bus width, and memory and
+/// SM clocks. Each device answers once.
+std::vector device_fingerprint() {
+  static std::mutex mutex;
+  static std::map> fingerprints;
+  int device = 0;
+  C10_CUDA_CHECK(cudaGetDevice(&device));
+  const std::lock_guard lock(mutex);
+  const auto known = fingerprints.find(device);
+  if (known != fingerprints.end()) {
+    return known->second;
+  }
+  const cudaDeviceAttr attributes[] = {
+      cudaDevAttrComputeCapabilityMajor,
+      cudaDevAttrComputeCapabilityMinor,
+      cudaDevAttrMultiProcessorCount,
+      cudaDevAttrL2CacheSize,
+      cudaDevAttrMaxRegistersPerMultiprocessor,
+      cudaDevAttrMaxSharedMemoryPerMultiprocessor,
+      cudaDevAttrMaxSharedMemoryPerBlockOptin,
+      cudaDevAttrMaxThreadsPerMultiProcessor,
+      cudaDevAttrMaxBlocksPerMultiprocessor,
+      cudaDevAttrGlobalMemoryBusWidth,
+      cudaDevAttrMemoryClockRate,
+      cudaDevAttrClockRate};
+  std::vector values;
+  for (const cudaDeviceAttr attribute : attributes) {
+    int value = 0;
+    C10_CUDA_CHECK(cudaDeviceGetAttribute(&value, attribute, device));
+    values.push_back(value);
+  }
+  std::vector fingerprint{10 * values[0] + values[1]};
+  fingerprint.insert(fingerprint.end(), values.begin() + 2, values.end());
+  TORCH_INTERNAL_ASSERT(fingerprint.size() == kFingerprintWords);
+  fingerprints.emplace(device, fingerprint);
+  return fingerprint;
+}
+
+/// The choice that ``table`` makes for a run of ``run_nodes`` nodes:
+/// kFittingPerLayer or the key of a fused shape where the measurement departed
+/// from the rules, and kFittingRules where it agreed with them or the table is
+/// empty, malformed or measured on a device of another fingerprint.
+int64_t table_choice(const std::vector& table, long run_nodes) {
+  constexpr size_t kHead = 2 + kFingerprintWords;
+  if (table.size() <= kHead || table[0] != kRouteTableMagic) {
+    return kFittingRules;
+  }
+  const std::vector fingerprint = device_fingerprint();
+  if (!std::equal(fingerprint.begin(), fingerprint.end(), table.begin() + 1)) {
+    return kFittingRules;
+  }
+  const int64_t intervals = table[kHead - 1];
+  if (intervals < 1 ||
+      table.size() != kHead + 2 * static_cast(intervals) - 1) {
+    return kFittingRules;
+  }
+  const int64_t* bounds = table.data() + kHead;
+  const int64_t* choices = bounds + intervals - 1;
+  for (int64_t i = 1; i < intervals - 1; ++i) {
+    if (bounds[i - 1] >= bounds[i]) {
+      return kFittingRules;
+    }
+  }
+  int64_t interval = 0;
+  while (interval < intervals - 1 && run_nodes > bounds[interval]) {
+    ++interval;
+  }
+  return choices[interval];
+}
+
+/// Row pitch of the transposed weight ``w``: its input width rounded up to
+/// whole quads.
+long transposed_pitch(const torch::Tensor& w) {
+  return (w.size(0) + 3) / 4 * 4;
+}
+
+/// Point the transposed weights of ``p`` into ``wt``, one (width, pitch)
+/// sheet per layer in order.
+void place_transposed_weights(FusedNetwork& p,
+                              const std::vector& ws,
+                              const float* wt) {
+  long offset = 0;
+  for (size_t l = 0; l < ws.size(); ++l) {
+    p.pitch[l] = static_cast(transposed_pitch(ws[l]));
+    p.wt[l] = wt + offset;
+    offset += ws[l].size(1) * p.pitch[l];
+  }
+}
+
+/// The plan of a fused network holds one layer per weight, at most
+/// ``kFusedLayers`` of them, which take at least one input feature and share
+/// one hidden width, the width that selects the kernel.
+void check_fused_plan(const FittingLayerPlan& plan,
+                      const std::vector& ws) {
+  TORCH_CHECK(plan.n_layer == static_cast(ws.size()) &&
+                  plan.n_layer <= kFusedLayers,
+              "graph_fitting: the fused network holds one plan layer per "
+              "weight and at most ",
+              kFusedLayers, " layers, got ", plan.n_layer, " and ", ws.size());
+  TORCH_CHECK(ws[0].size(0) > 0,
+              "graph_fitting: the fused network takes at least one input "
+              "feature");
+  for (int l = 1; l < plan.n_layer; ++l) {
+    TORCH_CHECK(ws[l].size(1) == ws[0].size(1),
+                "graph_fitting: the fused network takes one hidden width, got ",
+                ws[0].size(1), " and ", ws[l].size(1));
+  }
+}
+
+/// Whether the fused kernel computes the network: one to eight hidden layers
+/// of one width, every layer after the first square, and an input of at least
+/// one feature.
+bool fused_network_eligible(const FittingLayerPlan& plan,
+                            const std::vector& ws) {
+  if (plan.n_layer < 1 || plan.n_layer > kFusedLayers || ws[0].size(0) < 1) {
+    return false;
+  }
+  const long width = ws[0].size(1);
+  for (int l = 0; l < plan.n_layer; ++l) {
+    if (ws[l].size(1) != width || (l > 0 && ws[l].size(0) != width)) {
+      return false;
     }
   }
+  return true;
+}
+
+}  // namespace
+
+bool fitting_fused_serves(const FittingLayerPlan& plan,
+                          const std::vector& ws,
+                          long run_nodes) {
+  if (!fused_network_eligible(plan, ws)) {
+    return false;
+  }
+  const long width = ws[0].size(1);
+  const auto* properties = at::cuda::getCurrentDeviceProperties();
+  bool serves = false;
+  dispatch_fused_shape(width, run_nodes, [&](auto shape) {
+    using S = decltype(shape);
+    // The kernel streams its weights by asynchronous copies, which an image
+    // compiled for compute capability 7.x performs as plain loads and stores.
+    // The panels of a throughput shape share a multiprocessor and hide the
+    // weight stream behind each other, so they serve images of compute
+    // capability 7.0 and later; a latency shape runs each panel alone on its
+    // multiprocessor and keeps the stream ahead of its FMAs through the
+    // asynchronous copies of its ring alone, so it serves images of 8.0 and
+    // later. An image of an earlier architecture leaves the run to the
+    // per-layer kernels at every width.
+    //
+    // A run of few panels leaves each panel's pass through the whole network
+    // on one multiprocessor, while the per-layer kernels spread every layer
+    // over the device; a cluster splits each panel between several CTAs
+    // instead. A single CTA per panel serves once the run's 32-node panels
+    // cover the share of the multiprocessors its shape names -- half of them,
+    // and three quarters for a width of 256 -- from which the latency shapes
+    // outrun the per-layer kernels on sm_120.
+    const long panels = ceil_div(run_nodes, 32);
+    const long sms = properties->multiProcessorCount;
+    serves = fused_capacity().version >= (panels <= sms ? 80 : 70) &&
+             (S::cluster > 1 || 4 * panels >= S::serve_quarters * sms);
+  });
+  // graph_fitting_energy_gradient decides once for the runs of a tiled call
+  // and launches its tail with a shape of the tail's own size, so a run is
+  // served only where every shorter run finds a shape as well. The fused
+  // energy-force operator of DPA4C decides its tail by its own size but takes
+  // the same rule, so that the two operators route a run of one size alike. A
+  // tail of more panels than multiprocessors takes the throughput list of its
+  // width, as the served run does: one shape, or at a width of 96 the 64-node
+  // panels only where the device keeps them resident and otherwise the
+  // 32-node panels, which take less shared memory. A shorter tail takes the
+  // latency list, which ends with a single-CTA shape that launch_first_fit
+  // takes at any panel count once the device admits its shared memory; a
+  // device that admits less than every shape, below compute capability 8.0,
+  // launches no clusters, so its latency list serves every run size or none,
+  // and a run of one panel stands for every tail.
+  return serves && dispatch_fused_shape(width, 1, [](auto) {});
+}
+
+int64_t fitting_route(const FittingLayerPlan& plan,
+                      const std::vector& ws,
+                      long run_nodes,
+                      const std::vector& table) {
+  const int64_t choice = table_choice(table, run_nodes);
+  if (choice == kFittingPerLayer) {
+    return kFittingPerLayer;
+  }
+  if (choice > 0 && fused_network_eligible(plan, ws) &&
+      visit_fused_shapes(ws[0].size(1), [&](auto shape) {
+        using S = decltype(shape);
+        return S::key == choice && fused_shape_runs(run_nodes);
+      })) {
+    return choice;
+  }
+  return fitting_fused_serves(plan, ws, run_nodes) ? kFittingRules
+                                                   : kFittingPerLayer;
+}
+
+long fitting_fused_transposed_floats(const std::vector& ws) {
+  long total = 0;
+  for (const auto& w : ws) {
+    total += w.size(1) * transposed_pitch(w);
+  }
+  return total;
+}
+
+void fitting_fused_transpose(cudaStream_t stream,
+                             const FittingLayerPlan& plan,
+                             const std::vector& ws,
+                             float* wt) {
+  check_weights(ws);
+  check_fused_plan(plan, ws);
+  FusedNetwork p{};
+  TransposedSheets sheets{};
+  for (int l = 0; l < plan.n_layer; ++l) {
+    p.w[l] = ws[l].data_ptr();
+  }
+  p.din = (int)ws[0].size(0);
+  place_transposed_weights(p, ws, wt);
+  // The sheets lie in the caller's writable ``wt``.
+  int pitch = 0;
+  for (int l = 0; l < plan.n_layer; ++l) {
+    sheets.wt[l] = const_cast(p.wt[l]);
+    pitch = std::max(pitch, p.pitch[l]);
+  }
+  const int width = (int)ws[0].size(1);
+  const dim3 grid((unsigned)ceil_div(pitch, kTransposeTile),
+                  (unsigned)ceil_div(width, kTransposeTile),
+                  (unsigned)plan.n_layer);
+  transpose_weights_kernel<<>>(
+      p, width, sheets);
+  FITTING_CHECK_LAUNCH("graph_fitting weight transpose");
+}
+
+void fitting_energy_gradient_range(cudaStream_t stream,
+                                   const FittingLayerPlan& plan,
+                                   const float* x,
+                                   long input_width,
+                                   const long* atype,
+                                   const std::vector& ws,
+                                   const std::vector& bs,
+                                   const std::vector& resnets,
+                                   const torch::Tensor& w_head,
+                                   const torch::Tensor& b_head,
+                                   const torch::Tensor& bias_atom_e,
+                                   int64_t act,
+                                   const float* wt,
+                                   const double* d_e,
+                                   long run_nodes,
+                                   float* saved,
+                                   double* e,
+                                   float* d_x,
+                                   int64_t route) {
+  check_energy_parameters(ws, bs, resnets, w_head, b_head, bias_atom_e);
+  check_fused_plan(plan, ws);
+  FusedNetwork p{};
+  p.x = x;
+  p.atype = atype;
+  for (int l = 0; l < plan.n_layer; ++l) {
+    p.w[l] = ws[l].data_ptr();
+    p.b[l] = bs[l].numel() ? bs[l].data_ptr() : nullptr;
+    p.residual[l] = resnets[l] && ws[l].size(0) == ws[l].size(1) ? 1 : 0;
+  }
+  p.w_head = w_head.data_ptr();
+  p.b_head = b_head.numel() ? b_head.data_ptr() : nullptr;
+  p.bias_atom_e = bias_atom_e.data_ptr();
+  p.d_e = d_e;
+  p.saved = saved;
+  p.e = e;
+  p.d_x = d_x;
+  p.nodes = run_nodes;
+  p.din = (int)input_width;
+  p.n_layer = plan.n_layer;
+  place_transposed_weights(p, ws, wt);
+  const auto launch = [&](auto shape) {
+    using S = decltype(shape);
+    dispatch_activation(act, [&](auto tag) {
+      constexpr int kAct = decltype(tag)::value;
+      configure_fused_kernel();
+      cudaLaunchAttribute cluster;
+      const cudaLaunchConfig_t config = fused_launch_config(
+          ceil_div(run_nodes, S::rows), stream, &cluster);
+      C10_CUDA_CHECK(
+          cudaLaunchKernelEx(&config, fused_network_kernel, p));
+    });
+  };
+  const long width = ws[0].size(1);
+  const bool launched = route == kFittingRules
+                            ? dispatch_fused_shape(width, run_nodes, launch)
+                            : visit_fused_shapes(width, [&](auto shape) {
+                                if (decltype(shape)::key != route) {
+                                  return false;
+                                }
+                                launch(shape);
+                                return true;
+                              });
+  TORCH_CHECK(launched, "graph_fitting fused network: no fused shape serves ",
+              run_nodes, " nodes of hidden width ", width, " on route ", route,
+              "; fitting_fused_serves and fitting_route select the routes");
 }
 
 // Forward: per-atom energy (fp64 (N, 1)) plus the flat saved buffer of the
@@ -427,8 +2982,9 @@ std::tuple graph_fitting(
   }
   // Two-slot ping-pong for the activations: layer l writes slot ``l & 1``
   // while reading the previous layer's slot, so an activation is overwritten
-  // only after the next GEMM has consumed it (kernels run in stream order). A
-  // single-layer network requires only its output slot.
+  // only after the next GEMM has consumed it (kernels run in stream order).
+  // The last layer stores no activation; its slot holds the head partial dots,
+  // so a single-layer network requires only that slot.
   const int slots = plan.n_layer > 1 ? 2 : 1;
   auto act_buf = torch::empty({slots, n_node, plan.width_max}, f32);
   float* activation[2] = {
@@ -471,13 +3027,17 @@ void graph_fitting_backward_core(torch::Tensor d_e,
   const c10::cuda::CUDAGuard device_guard(saved.device());
   auto f32 = saved.options().dtype(torch::kFloat32);
   auto d_e_c = d_e.contiguous();
-  auto dh = torch::empty({n_node, plan.width_max}, f32);
-  auto dh_next = plan.n_layer > 1 ? torch::empty({n_node, plan.width_max}, f32)
-                                  : torch::empty({0}, f32);
+  const auto scratch = [&](bool needed) {
+    return needed ? torch::empty({n_node, plan.width_max}, f32)
+                  : torch::empty({0}, f32);
+  };
+  auto dh = scratch(plan.n_layer > 1);
+  auto dh_next = scratch(plan.n_layer > 2);
   fitting_backward_range(at::cuda::getCurrentCUDAStream(), plan,
                          d_e_c.data_ptr(), saved.data_ptr(), ws,
-                         bs, resnets, w_head, act, n_node, dh.data_ptr(),
-                         plan.n_layer > 1 ? dh_next.data_ptr() : nullptr,
+                         bs, resnets, w_head, act, n_node,
+                         plan.n_layer > 1 ? dh.data_ptr() : nullptr,
+                         plan.n_layer > 2 ? dh_next.data_ptr() : nullptr,
                          d_x.data_ptr());
 }
 
@@ -509,7 +3069,9 @@ torch::Tensor graph_fitting_backward(torch::Tensor d_e,
 // node, so no tile needs anything from another and nothing is recomputed.
 //
 // ``tile`` of zero, or any value covering the system, degenerates to a single
-// run and reproduces the untiled sequence exactly.
+// run and reproduces the untiled sequence exactly. A network and run that the
+// fused network serves take it instead of the per-layer sequence, with the
+// same results bit for bit.
 //
 // The cotangent replaces the descriptor in place. A tile reads its descriptor
 // rows in the first layer and writes their cotangent in the last, and no later
@@ -545,17 +3107,31 @@ torch::Tensor graph_fitting_energy_gradient(torch::Tensor x,
       "per node on the device of x");
 
   const long run = tile > 0 ? std::min(tile, n_node) : n_node;
-  const int slots = plan.n_layer > 1 ? 2 : 1;
   auto saved = torch::empty({run * plan.saved_width()}, f32);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (fitting_fused_serves(plan, ws, run)) {
+    auto transposed = torch::empty({fitting_fused_transposed_floats(ws)}, f32);
+    fitting_fused_transpose(stream, plan, ws, transposed.data_ptr());
+    for (long begin = 0; begin < n_node; begin += run) {
+      const long count = std::min(run, n_node - begin);
+      fitting_energy_gradient_range(
+          stream, plan, x.data_ptr() + begin * input_width, input_width,
+          atype.data_ptr() + begin, ws, bs, resnets, w_head, b_head,
+          bias_atom_e, act, transposed.data_ptr(),
+          seed_c.data_ptr() + begin, count, saved.data_ptr(),
+          e.data_ptr() + begin,
+          x.data_ptr() + begin * input_width, kFittingRules);
+    }
+    return e;
+  }
   // The backward reads only the saved pre-activations, never the activations,
   // so the forward ping-pong and the two cotangent buffers never hold live
   // data at the same time and share one allocation.
+  const int slots = plan.n_layer > 1 ? 2 : 1;
   auto scratch = torch::empty({slots, run, plan.width_max}, f32);
   float* slot[2] = {
       scratch[0].data_ptr(),
       slots > 1 ? scratch[1].data_ptr() : scratch[0].data_ptr()};
-
-  auto stream = at::cuda::getCurrentCUDAStream();
   for (long begin = 0; begin < n_node; begin += run) {
     const long count = std::min(run, n_node - begin);
     fitting_forward_range(
@@ -572,9 +3148,169 @@ torch::Tensor graph_fitting_energy_gradient(torch::Tensor x,
   return e;
 }
 
+namespace {
+
+// Queries of the freeze-time measurement of the fitting routes
+// (deepmd/pt_expt/kernels/fitting_routes.py) and of its tests. They answer for
+// CUDA device ``device``.
+
+/// The fingerprint that a fitting-route table measured on the device carries.
+std::vector graph_fitting_route_fingerprint(int64_t device) {
+  const c10::cuda::CUDAGuard device_guard(
+      static_cast(device));
+  return device_fingerprint();
+}
+
+/// The keys of the fused shapes of hidden width ``width`` that a measurement
+/// compares on a run of ``run_nodes`` nodes (visit_compared_shapes) and that
+/// the device runs on it, in the order of the dispatch.
+std::vector graph_fitting_route_choices(int64_t device,
+                                                 int64_t width,
+                                                 int64_t run_nodes) {
+  const c10::cuda::CUDAGuard device_guard(
+      static_cast(device));
+  const long sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  std::vector keys;
+  visit_compared_shapes(width, ceil_div(run_nodes, 32) <= sms, [&](auto shape) {
+    using S = decltype(shape);
+    if (fused_shape_runs(run_nodes)) {
+      keys.push_back(S::key);
+    }
+    return false;
+  });
+  return keys;
+}
+
+/// The route that the rules take for the network ``ws`` on a run of
+/// ``run_nodes`` nodes: kFittingPerLayer or the key of the fused shape the
+/// dispatch selects.
+int64_t graph_fitting_route_rule(std::vector ws,
+                                 int64_t run_nodes) {
+  TORCH_CHECK(!ws.empty() && ws[0].is_cuda(),
+              "graph_fitting_route_rule: the weights must be CUDA tensors");
+  const c10::cuda::CUDAGuard device_guard(ws[0].device());
+  const FittingLayerPlan plan = fitting_layer_plan(ws);
+  if (!fitting_fused_serves(plan, ws, run_nodes)) {
+    return kFittingPerLayer;
+  }
+  int64_t key = kFittingPerLayer;
+  dispatch_fused_shape(ws[0].size(1), run_nodes,
+                       [&](auto shape) { key = decltype(shape)::key; });
+  return key;
+}
+
+/// The route that the fused energy-force operator of DPA4C takes for the
+/// network ``ws`` on a run of ``run_nodes`` nodes under the fitting-route
+/// table ``table`` (fitting_route): kFittingPerLayer, kFittingRules or the key
+/// of a fused shape.
+int64_t graph_fitting_route_resolve(std::vector ws,
+                                    int64_t run_nodes,
+                                    std::vector table) {
+  TORCH_CHECK(!ws.empty() && ws[0].is_cuda(),
+              "graph_fitting_route_resolve: the weights must be CUDA tensors");
+  const c10::cuda::CUDAGuard device_guard(ws[0].device());
+  return fitting_route(fitting_layer_plan(ws), ws, run_nodes, table);
+}
+
+/// Inclusive upper bounds of the run sizes, up to ``max_nodes`` and ending
+/// with it, between which none of the routes that a measurement compares
+/// (visit_compared_shapes and the per-layer kernels) changes the number of
+/// waves of its grids on the device, and no rule changes the route of the
+/// network of hidden width ``width`` and input width ``din``: the cluster
+/// capacities and the serve gates of the latency shapes, the waves of the
+/// single-CTA shapes compared above one panel per multiprocessor, the waves of
+/// the per-layer tiles on every column count of the network and the switch
+/// between the two tiles, and the switch between the latency and the
+/// throughput shapes. A route keeps its waves inside such a band, but not the
+/// fill of its last wave, which the measurement samples at both ends of the
+/// band.
+std::vector graph_fitting_route_bands(int64_t device,
+                                               int64_t width,
+                                               int64_t din,
+                                               int64_t max_nodes) {
+  const c10::cuda::CUDAGuard device_guard(
+      static_cast(device));
+  const int64_t sms =
+      at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  const int64_t latency_nodes = 32 * sms;
+  std::set bounds;
+  const auto add = [&](int64_t nodes) {
+    if (nodes >= 1 && nodes < max_nodes) {
+      bounds.insert(nodes);
+    }
+  };
+  add(latency_nodes);
+  visit_compared_shapes(width, true, [&](auto shape) {
+    using S = decltype(shape);
+    const FusedCapacity capacity = fused_capacity();
+    if constexpr (S::cluster > 1) {
+      add(std::min(capacity.clusters * S::rows, latency_nodes));
+    } else {
+      add(32 * (ceil_div(S::serve_quarters * sms, 4) - 1));
+    }
+    return false;
+  });
+  visit_compared_shapes(width, false, [&](auto shape) {
+    using S = decltype(shape);
+    const int64_t wave =
+        static_cast(fused_capacity().resident) * sms * S::rows;
+    for (int64_t nodes = wave; wave > 0 && nodes < max_nodes; nodes += wave) {
+      if (nodes > latency_nodes) {
+        add(nodes);
+      }
+    }
+    return false;
+  });
+  // The per-layer kernels take the small tile up to the run size at which its
+  // grid reaches four and a half CTAs per multiprocessor (dispatch_tile), and
+  // the large tile beyond it, on each GEMM of the network.
+  for (const int64_t columns : {width, din}) {
+    const int64_t small_limit =
+        SmallTile::rows * (9 * sms / (2 * ceil_div(columns, SmallTile::cols)));
+    add(small_limit);
+    const auto add_tile_waves = [&](auto tile, int64_t from, int64_t to) {
+      using Tile = decltype(tile);
+      const int64_t wave = static_cast(tile_resident()) * sms;
+      const int64_t tiles = ceil_div(columns, Tile::cols);
+      for (int64_t waves = 1; wave > 0; ++waves) {
+        const int64_t nodes = Tile::rows * (waves * wave / tiles);
+        if (nodes >= to) {
+          break;
+        }
+        if (nodes > from) {
+          add(nodes);
+        }
+      }
+    };
+    add_tile_waves(SmallTile{}, 0, small_limit);
+    add_tile_waves(LargeTile{}, small_limit, max_nodes);
+  }
+  bounds.insert(max_nodes);
+  return {bounds.begin(), bounds.end()};
+}
+
+}  // namespace
+
 TORCH_LIBRARY_FRAGMENT(deepmd, m) {
   m.impl("graph_fitting", torch::kCUDA, &graph_fitting);
   m.impl("graph_fitting_backward", torch::kCUDA, &graph_fitting_backward);
   m.impl("graph_fitting_energy_gradient", torch::kCUDA,
          &graph_fitting_energy_gradient);
+  m.def("graph_fitting_route_fingerprint(int device) -> int[]",
+        &graph_fitting_route_fingerprint);
+  m.def(
+      "graph_fitting_route_choices(int device, int width, int run_nodes) "
+      "-> int[]",
+      &graph_fitting_route_choices);
+  m.def(
+      "graph_fitting_route_bands(int device, int width, int din, "
+      "int max_nodes) -> int[]",
+      &graph_fitting_route_bands);
+  m.def("graph_fitting_route_rule(Tensor[] ws, int run_nodes) -> int");
+  m.impl("graph_fitting_route_rule", torch::kCUDA, &graph_fitting_route_rule);
+  m.def(
+      "graph_fitting_route_resolve(Tensor[] ws, int run_nodes, int[] table) "
+      "-> int");
+  m.impl("graph_fitting_route_resolve", torch::kCUDA,
+         &graph_fitting_route_resolve);
 }
diff --git a/source/op/pt/graph_ops.h b/source/op/pt/graph_ops.h
index 90ee0899a0..8e8540b45b 100644
--- a/source/op/pt/graph_ops.h
+++ b/source/op/pt/graph_ops.h
@@ -1,12 +1,14 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 //
-// Entry points of the fused graph-lower operator suite, shared so the fused
-// energy-force operator (dpa1_graph_energy_force.cu) can drive the forward and
-// backward passes directly rather than through the dispatcher. The definitions
-// live in dpa1_graph_descriptor.cu, graph_fitting.cu and edge_force_virial.cu;
-// the operator schemas registered from those files are the public interface.
-// The compressed descriptor (dpa1_graph_compress.cu) is not part of the
-// end-to-end operator, so its entries stay private to that translation unit.
+// Entry points of the fused graph-lower operator suite, shared so that the
+// fused energy-force operators -- that of DPA1 (dpa1_graph_energy_force.cu) and
+// the compressed DPA4C one (dpa4c/graph_compress.cu) -- drive the descriptor,
+// fitting and force passes directly rather than through the dispatcher. The
+// definitions live in dpa1_graph_descriptor.cu, graph_fitting.cu and
+// edge_force_virial.cu; the operator schemas registered from those files are
+// the public interface. The compressed DPA1 descriptor (dpa1_graph_compress.cu)
+// is not part of an end-to-end operator, so its entries stay private to that
+// translation unit.
 
 #pragma once
 
@@ -123,7 +125,12 @@ torch::Tensor graph_fitting_backward(torch::Tensor d_e,
 // Evaluate the fitting network over one contiguous run of nodes. Every
 // full-width pointer is already indexed from the run's first node, so the same
 // code serves the whole node axis and a single tile of it. ``saved`` and
-// ``activation`` are sized for the run rather than for the system.
+// ``activation`` are sized for the run rather than for the system. The
+// kernels address the hidden rows of ``saved``, ``activation`` and, in the
+// backward, ``dh`` and ``dh_next`` in aligned float4 quads, so these buffers
+// are 16-byte aligned, as every torch allocation is; ``x`` and ``d_x`` may lie
+// at any float boundary. The weights, biases and head weight are contiguous
+// 16-byte aligned fp32 tensors on the current device (check_weights).
 void fitting_forward_range(cudaStream_t stream,
                            const FittingLayerPlan& plan,
                            const float* x,
@@ -156,6 +163,115 @@ void fitting_backward_range(cudaStream_t stream,
                             float* dh_next,
                             float* d_x);
 
+// Whether the fused fitting kernel serves a run of ``run_nodes`` nodes: an
+// input of at least one feature, one to eight hidden layers of one width it is
+// compiled for, every layer after the first square, and shapes whose shared
+// memory the device admits for the run and for a run of one panel. A run of
+// more 32-node panels than multiprocessors needs a kernel image compiled for
+// compute_70 or later, and at a width that 128 divides it is left to the
+// per-layer kernels where the fp32 ridge of the device lies below width / 4. A
+// run of at most one panel per multiprocessor needs an image compiled for
+// compute_80 or later and panels that either split over clusters of CTAs the
+// device keeps resident at once or, one CTA each, cover half of the
+// multiprocessors (three quarters at a width of 256). At a width of 384 it
+// serves runs of at most one panel per multiprocessor.
+bool fitting_fused_serves(const FittingLayerPlan& plan,
+                          const std::vector& ws,
+                          long run_nodes);
+
+// Routes of the fitting network that fitting_route returns besides the key of
+// a fused shape, which is positive.
+//
+// kFittingPerLayer: the per-layer kernels, fitting_forward_range followed by
+// fitting_backward_range.
+constexpr int64_t kFittingPerLayer = -1;
+// kFittingRules: the fused kernel, of the shape that the rules of the dispatch
+// select for the run size; in a fitting-route table, the choice that leaves
+// the run to the rules.
+constexpr int64_t kFittingRules = 0;
+
+// Route of the fitting network over a run of ``run_nodes`` nodes. ``table`` is
+// a fitting-route table measured at freeze time on the device of the run
+// (graph_fitting.cu), which names a route only where the measurement departs
+// from the rules and names the rules elsewhere. The run takes the route that
+// the table names for its size, or the rules where the table names them, is
+// empty, was measured on another device, or names a fused shape that the
+// device cannot run on the run size or that is not compiled for the network.
+// The rules route the run to the fused kernel where fitting_fused_serves holds
+// and to the per-layer kernels otherwise. Every route evaluates the same
+// contractions in the same order, so the route sets only the speed.
+int64_t fitting_route(const FittingLayerPlan& plan,
+                      const std::vector& ws,
+                      long run_nodes,
+                      const std::vector& table);
+
+// Floats of the transposed weights the fused fitting kernel streams in its
+// backward, and their layout from the weights: one (width, din rounded up to
+// four) row-major sheet per layer, zero past din. Both serve a network for
+// which fitting_fused_serves holds; the transpose checks that the network has
+// one to eight layers of one hidden width and at least one input feature.
+long fitting_fused_transposed_floats(const std::vector& ws);
+void fitting_fused_transpose(cudaStream_t stream,
+                             const FittingLayerPlan& plan,
+                             const std::vector& ws,
+                             float* wt);
+
+// Energy and input cotangent of one run of nodes in a single launch: the
+// forward, the head and the backward from the energy cotangent ``d_e``. The
+// results are bitwise those of fitting_forward_range followed by
+// fitting_backward_range. ``saved`` holds n_layer sheets of (run_nodes,
+// width): the pre-activations of every layer but the last, then the sheet
+// through which the CTAs of a cluster exchange their column slices; like
+// ``wt``, it is 16-byte aligned. ``wt`` holds the sheets that
+// fitting_fused_transpose writes for the same weights, and
+// fitting_fused_serves holds for the network and for a run of ``run_nodes``
+// or more, or ``route`` is a fused route that fitting_route returned for a run
+// of ``run_nodes``. ``d_x`` may alias ``x``.
+void fitting_energy_gradient_range(cudaStream_t stream,
+                                   const FittingLayerPlan& plan,
+                                   const float* x,
+                                   long input_width,
+                                   const long* atype,
+                                   const std::vector& ws,
+                                   const std::vector& bs,
+                                   const std::vector& resnets,
+                                   const torch::Tensor& w_head,
+                                   const torch::Tensor& b_head,
+                                   const torch::Tensor& bias_atom_e,
+                                   int64_t act,
+                                   const float* wt,
+                                   const double* d_e,
+                                   long run_nodes,
+                                   float* saved,
+                                   double* e,
+                                   float* d_x,
+                                   int64_t route);
+
+// Source-attributed per-node virial of a canonical graph from stored edge
+// cotangents, atom_virial[node] = sum(src=node) -g_e (x) edge_vec_e. The edge
+// tensors are contiguous fp32 of shape (S, 3), the source CSR holds
+// ``node_count + 1`` int64 row pointers and S int64 positions, and every
+// tensor resides on the current CUDA device. Returns (node_count, 3, 3) in
+// fp32.
+torch::Tensor canonical_source_virial(const torch::Tensor& edge_gradient,
+                                      const torch::Tensor& edge_vec,
+                                      const torch::Tensor& source_row_ptr,
+                                      const torch::Tensor& source_order,
+                                      long node_count);
+
+// Exclusive prefix of the per-frame node counts as int64, so that frame ``f``
+// spans ``[row_ptr[f], row_ptr[f + 1])`` of the node axis.
+torch::Tensor frame_row_pointer(const torch::Tensor& n_node_per_frame);
+
+// Per-frame sum of an fp32 per-node virial with shape (N, 9), accumulated and
+// returned in fp64 with shape (nf, 3, 3).
+torch::Tensor frame_virial_sum(const torch::Tensor& node_virial,
+                               const torch::Tensor& n_node_per_frame);
+
+// Per-frame total of a node scalar with shape (N, 1), with shape (nf, 1).
+torch::Tensor frame_scalar_sum(torch::Tensor node_scalar,
+                               torch::Tensor n_node_per_frame);
+
 // Scatter dE/d(edge_vec) into per-node force, per-frame virial and (optional)
 // per-node virial. A rank-two ``edge_spin_gradient`` adds the per-source total
 // of the magnetic cotangent, which shares the source grouping the force
diff --git a/source/op/pt/long_range/charge_sum.h b/source/op/pt/long_range/charge_sum.h
new file mode 100644
index 0000000000..dcccb69393
--- /dev/null
+++ b/source/op/pt/long_range/charge_sum.h
@@ -0,0 +1,94 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// The two evaluations behind deepmd::latent_charge_sum (schema and formulas in
+// ops.cc): the direct sum over the pairs of non-periodic frames
+// (direct_sum.cu) and the smooth particle-mesh Ewald sum of periodic frames
+// (mesh_sum.cu). Both take float64 positions and float32 raw charges of the
+// nodes of the frames on one CUDA device, apply the constraint, and evaluate
+// the owner nodes [owner_begin, owner_begin + owner_count) against every node
+// of their frames. Owners other than every node lie in a single frame, whose
+// nodes are then the charged nodes of every domain of ``communicator``, over
+// which the per-frame sums of the owners are added.
+
+#pragma once
+
+#include 
+
+#include 
+#include 
+
+#include "domain_exchange.h"
+
+namespace deepmd {
+namespace long_range {
+
+// The channel count and the Gaussian count that the kernels are instantiated
+// for; the Gaussians of a kernel are staged in shared memory.
+constexpr int kMaxChannels = 8;
+constexpr int kMaxGaussians = 32;
+
+// The interaction kernel: erf(alpha r) / r (LES) or a sum of Gaussians (SOG).
+enum KernelKind : int { kLes = 0, kSog = 1 };
+
+// The atomic energy (owners, 1) float64, the charge gradient (owners, C), the
+// force (owners, 3), the atomic virial (owners, 9) and the virial of the
+// owners of every frame (F, 9), float32.
+using ChargeSumOutputs = std::tuple;
+
+// The direct sum, of at least one node and one frame. ``coefficient`` holds
+// the float64 coefficients of the kernel on the device.
+ChargeSumOutputs direct_sum(const torch::Tensor& coord,
+                            const torch::Tensor& charge,
+                            const torch::Tensor& node_mask,
+                            const torch::Tensor& n_node_per_frame,
+                            const torch::Tensor& total_charge,
+                            const torch::Tensor& coefficient,
+                            int64_t kernel,
+                            int64_t channel,
+                            double prefactor,
+                            int owner_begin,
+                            int owner_count,
+                            const torch::Tensor* communicator);
+
+// The mesh sum, of at least one node and one frame. ``box`` holds the cell of
+// every frame, float64 on the host;
+// ``coefficient`` the float64 coefficients of the kernel on the device.
+ChargeSumOutputs mesh_sum(const torch::Tensor& coord,
+                          const torch::Tensor& charge,
+                          const torch::Tensor& node_mask,
+                          const torch::Tensor& n_node_per_frame,
+                          const torch::Tensor& total_charge,
+                          const torch::Tensor& box,
+                          const torch::Tensor& coefficient,
+                          int64_t kernel,
+                          int64_t channel,
+                          double prefactor,
+                          int64_t order,
+                          double spacing,
+                          int owner_begin,
+                          int owner_count,
+                          const torch::Tensor* communicator);
+
+// deepmd::latent_charge_sum (schema and formulas in ops.cc): validates the
+// inputs and evaluates the frames by the direct or the mesh sum, across the
+// domains of ``communicator`` when it holds several ranks.
+ChargeSumOutputs latent_charge_sum(torch::Tensor coord,
+                                   torch::Tensor charge,
+                                   torch::Tensor node_mask,
+                                   torch::Tensor n_node_per_frame,
+                                   torch::Tensor total_charge,
+                                   c10::optional box,
+                                   torch::Tensor coefficient,
+                                   torch::Tensor communicator,
+                                   int64_t kernel,
+                                   int64_t channel,
+                                   double prefactor,
+                                   int64_t order,
+                                   double spacing);
+
+}  // namespace long_range
+}  // namespace deepmd
diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu
new file mode 100644
index 0000000000..677f2e25d7
--- /dev/null
+++ b/source/op/pt/long_range/direct_sum.cu
@@ -0,0 +1,1037 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// CUDA kernels of the latent-charge long range of non-periodic frames, by the
+// direct sum over their pairs (schema and formulas in ops.cc, the entry in
+// latent_charge_sum.cu).
+//
+// Every node reads every node of its frame, so the memory is linear in the
+// node count while the work is quadratic. One thread owns one node and walks
+// the nodes of its frame in tiles staged in shared memory; the partner of a
+// thread at a given step is the same for the whole block, so each tile entry
+// is a broadcast read. Each pair is evaluated by both of its nodes, which
+// keeps every sum in the registers of its owner: no atomic operation is
+// needed, and the result does not depend on scheduling.
+//
+// Positions are carried relative to the first node of their frame as a
+// double-float pair: a float32 head and the float32 tail of the remainder.
+// The displacement of two nodes, formed as the difference of the heads plus
+// the difference of the tails, keeps the float32 precision of the
+// displacement itself wherever the frame lies, while every pair evaluates in
+// float32; the far field of the erf kernel forms it from the heads alone (see
+// far_batch). For the erf kernel the positions are scaled by alpha, so that
+// the pair loop evaluates erf(x) / x and its slope without a coefficient, and
+// the reduction restores the units.
+//
+// A thread sums one tile in float32 and adds the tile sums with compensated
+// float32 arithmetic. The partner axis is split into contiguous chunks, so
+// that a launch fills the device several times over; every chunk writes its
+// sums, and a second pass adds the chunks in a fixed order in float64. A
+// system too small to fill the device with tiles of kWideTile nodes takes
+// tiles of one warp, so that its pairs spread over more multiprocessors. Its
+// chunks hold kNarrowChunkTiles tiles each, counted from the first partner of
+// each block, rather than equal parts of the partner range: the chunks of a
+// block then do not depend on the partners past the frames of its nodes.
+//
+// The charge constraint and the per-frame virial reduce over the nodes of
+// every frame, in slices whose float64 sums are added in a fixed order, so
+// that every output is independent of scheduling.
+//
+// The nodes whose sums a call evaluates, its owners, are either every node
+// or, across the domains of a decomposed system, the charged nodes of one
+// MPI rank, which the entry gathers with those of every rank into one frame;
+// each rank evaluates its own nodes against the whole frame, so that the
+// quadratic work divides among the ranks. Every rank shifts the constrained
+// channel by the same sum of the gathered charges; the mean potential of the
+// constraint adds the sums of the ranks in rank order, and the per-frame
+// virial is the share of the owners of the rank.
+
+#include 
+#include 
+#include 
+
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+
+#include "charge_sum.h"
+#include "frame_sums.cuh"
+
+namespace {
+
+// Threads per block of the per-node kernels.
+constexpr int kThreads = 128;
+// Nodes of a tile of the pair kernel: a block owns one tile of nodes and
+// walks the partners of its chunk in tiles of the same length. Every thread
+// sums the pairs of a tile in float32 and adds the tile sum to its
+// compensated sums.
+constexpr int kWideTile = 128;
+constexpr int kNarrowTile = 32;
+// Partners evaluated together, which share every coefficient load.
+constexpr int kBatch = 4;
+// Fills of the resident block slots of the device that one launch aims for.
+constexpr int kWaves = 8;
+// Tiles per chunk of wide tiles where a frame that spans every partner takes
+// the most chunks it may: pair_launch rounds the chunk count up from its tile
+// count over this.
+constexpr int kWideChunkTiles = 2;
+// Tiles of a chunk of narrow tiles: two, so that a chunk that holds a tile
+// entirely past the partners of a frame holds at most one tile with partners
+// of that frame. The compensated sum of one tile carries no error, so that
+// the zero sums of the tiles past the frame leave it unchanged, as the frame
+// on its own would.
+constexpr int kNarrowChunkTiles = 2;
+
+using deepmd::long_range::kLes;
+using deepmd::long_range::kMaxChannels;
+using deepmd::long_range::kMaxGaussians;
+using deepmd::long_range::kSog;
+
+// Sums of one node: the potential of every channel, the force direction and
+// the six independent components of the virial.
+template 
+constexpr int sum_width() {
+  return kChannels + 3 + 6;
+}
+
+// Running sum with Kahan compensation.
+struct CompensatedSum {
+  float sum = 0.f;
+  float error = 0.f;
+
+  __device__ __forceinline__ void add(float value) {
+    const float corrected = value - error;
+    const float total = sum + corrected;
+    error = (total - sum) - corrected;
+    sum = total;
+  }
+};
+
+// Writes the double-float positions relative to the first node of each
+// frame, scaled by alpha for the erf kernel, the frame of every node, and the
+// charges of the pairs: the raw charges of the charged nodes, shifted in the
+// constrained channel, and zero for every other node. ``alpha`` is a device
+// scalar, or null for a kernel read in plain units; ``frame_stat`` is null
+// without a constraint.
+__global__ void prepare_nodes(const double* __restrict__ coord,
+                              const float* __restrict__ charge,
+                              const bool* __restrict__ node_mask,
+                              const int64_t* __restrict__ frame_ptr,
+                              int frame_count,
+                              int node_count,
+                              int channels,
+                              int channel,
+                              const double* __restrict__ alpha,
+                              const double* __restrict__ frame_stat,
+                              float4* __restrict__ head,
+                              float4* __restrict__ tail,
+                              int* __restrict__ frame,
+                              float* __restrict__ pair_charge) {
+  const int node = blockIdx.x * blockDim.x + threadIdx.x;
+  if (node >= node_count) {
+    return;
+  }
+  const double scale = alpha == nullptr ? 1.0 : *alpha;
+  const int f = frame_of(frame_ptr, frame_count, node);
+  double3 offset = make_double3(0.0, 0.0, 0.0);
+  if (f >= 0) {
+    const int64_t first = frame_ptr[f];
+    offset.x = scale * (coord[3 * node + 0] - coord[3 * first + 0]);
+    offset.y = scale * (coord[3 * node + 1] - coord[3 * first + 1]);
+    offset.z = scale * (coord[3 * node + 2] - coord[3 * first + 2]);
+  }
+  const float hx = __double2float_rn(offset.x);
+  const float hy = __double2float_rn(offset.y);
+  const float hz = __double2float_rn(offset.z);
+  head[node] = make_float4(hx, hy, hz, 0.f);
+  tail[node] = make_float4(__double2float_rn(offset.x - hx),
+                           __double2float_rn(offset.y - hy),
+                           __double2float_rn(offset.z - hz), 0.f);
+  frame[node] = f;
+  const bool charged = f >= 0 && node_mask[node];
+  const double shift = charged && frame_stat != nullptr
+                           ? frame_stat[f * kStatWidth + kStatShift]
+                           : 0.0;
+  for (int c = 0; c < channels; ++c) {
+    const int64_t index = static_cast(node) * channels + c;
+    const double value = charged ? static_cast(charge[index]) : 0.0;
+    pair_charge[index] =
+        static_cast(c == channel ? value - shift : value);
+  }
+}
+
+// K and K'(r) / r of a batch of pairs at the squared distances ``r2``.
+//
+// The erf kernel reads positions scaled by alpha, x = alpha r, and evaluates
+// erf(x) / x and its slope (2 / sqrt(pi) exp(-x^2) - erf(x) / x) / x^2, with
+// erf(x) = 1 - t P(t) exp(-x^2), t = 1 / (1 + p x), and the coefficients of
+// Abramowitz and Stegun 7.1.26 (absolute error below 1.5e-7), so that one
+// exponential serves the kernel and its slope. The two terms of the slope
+// agree to second order in x, which amplifies the error of the approximation
+// by 1 / x^2; below x^2 = kSeriesSquared the kernel and its slope therefore
+// come from their Taylor series in x^2,
+//
+//   erf(x) / x  = 2 / sqrt(pi) sum_n (-1)^n x^(2n) / (n! (2n + 1)),
+//   slope       = 2 / sqrt(pi) sum_{n >= 1} (-1)^n 2n x^(2n - 2) /
+//                 (n! (2n + 1)),
+//
+// whose first six terms reach float32 precision there. The relative error of
+// the kernel stays below 1e-6 for every x, that of its slope below 7e-6, the
+// largest just above the threshold. A warp evaluates the series only for a
+// batch with a pair below the threshold: a batch that pairs a node with
+// itself, or a contact shorter than 0.4 / alpha, which for an alpha of 0.5 per
+// angstrom or more is shorter than every bond but that of H2. Every other
+// batch evaluates the closed form without a branch.
+//
+// Beyond x^2 = kFarSquared the float32 kernel equals 1 / x and its slope
+// -1 / x^3 bit for bit: exp(-x^2) and erfc(x) fall below half a unit in the
+// last place of the terms they would correct. A warp whose batch lies there
+// entirely skips the evaluation of erf, which most pairs of a large system
+// do, and forms the displacements of the batch from the heads alone. The
+// omitted tails move a displacement by at most 2^-24 (|x_i| + |x_j|), with x
+// the alpha-scaled offset of a node from the first node of its frame: the
+// rounding of float32 positions over the extent of the frame. Against a
+// distance of at least sqrt(kFarSquared) = 4.5, the relative error of the
+// displacement stays below 2^-24 (|x_i| + |x_j|) / 4.5. For a cluster of
+// 100,000 carbon atoms at alpha = 0.5 per angstrom, with offsets of up to 26
+// units, the potential and the force keep the error of the exact
+// displacements.
+//
+// The sum of Gaussians evaluates sum_k a_k 2^(-c_k r^2) and its slope
+// sum_k s_k 2^(-c_k r^2), s_k = -2 ln(2) c_k a_k, from the staged
+// coefficients ``{a_k, -c_k, s_k}``.
+constexpr float kFarSquared = 20.25f;
+constexpr float kSeriesSquared = 0.16f;
+
+// A series in x^2 with the coefficients ``c``, in increasing powers, by
+// Horner's rule.
+template 
+__device__ __forceinline__ float even_series(const float (&c)[kTerms],
+                                             float x2) {
+  float sum = c[kTerms - 1];
+#pragma unroll
+  for (int n = kTerms - 2; n >= 0; --n) {
+    sum = fmaf(sum, x2, c[n]);
+  }
+  return sum;
+}
+
+// Whether every pair of the batch of every thread of the warp lies in the far
+// field of the erf kernel, judged from the head displacements, whose squared
+// lengths it writes to ``r2``. The sum of Gaussians has no far field.
+template 
+__device__ __forceinline__ bool far_batch(const float (&dx)[kBatch],
+                                          const float (&dy)[kBatch],
+                                          const float (&dz)[kBatch],
+                                          float (&r2)[kBatch]) {
+  if constexpr (kKernel == kSog) {
+    return false;
+  } else {
+    float nearest = std::numeric_limits::infinity();
+#pragma unroll
+    for (int b = 0; b < kBatch; ++b) {
+      r2[b] = fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b]));
+      nearest = fminf(nearest, r2[b]);
+    }
+    return __all_sync(0xffffffffu, nearest >= kFarSquared);
+  }
+}
+
+// erf(x) / x and its slope by the approximation of Abramowitz and Stegun at
+// the squared distances ``r2``.
+__device__ __forceinline__ void les_closed_form(const float (&r2)[kBatch],
+                                                float (&value)[kBatch],
+                                                float (&slope)[kBatch]) {
+#pragma unroll
+  for (int b = 0; b < kBatch; ++b) {
+    const float inverse = rsqrtf(r2[b]);
+    const float x = r2[b] * inverse;
+    const float t = __fdividef(1.f, fmaf(0.3275911f, x, 1.f));
+    const float poly =
+        t *
+        fmaf(
+            t,
+            fmaf(t, fmaf(t, fmaf(t, 1.061405429f, -1.453152027f), 1.421413741f),
+                 -0.284496736f),
+            0.254829592f);
+    const float gaussian = exp2f(-1.44269504088896341f * r2[b]);
+    value[b] = fmaf(-poly, gaussian, 1.f) * inverse;
+    slope[b] =
+        fmaf(1.12837916709551257f, gaussian, -value[b]) * (inverse * inverse);
+  }
+}
+
+template 
+__device__ __forceinline__ void evaluate_kernel(const float (&r2)[kBatch],
+                                                const float4* gaussians,
+                                                int gaussian_count,
+                                                float (&value)[kBatch],
+                                                float (&slope)[kBatch]) {
+  if constexpr (kKernel == kLes) {
+    float nearest = r2[0];
+#pragma unroll
+    for (int b = 1; b < kBatch; ++b) {
+      nearest = fminf(nearest, r2[b]);
+    }
+    if (!__any_sync(0xffffffffu, nearest < kSeriesSquared)) {
+      les_closed_form(r2, value, slope);
+      return;
+    }
+    constexpr float kValueSeries[6] = {1.128379167f,    -0.376126389f,
+                                       0.1128379167f,   -0.02686617065f,
+                                       0.005223977625f, -0.0008548327023f};
+    constexpr float kSlopeSeries[6] = {-0.7522527781f,   0.4513516668f,
+                                       -0.1611970239f,   0.041791821f,
+                                       -0.008548327023f, 0.001446639958f};
+    // The closed form reads the threshold where the series is selected.
+    float clamped[kBatch];
+#pragma unroll
+    for (int b = 0; b < kBatch; ++b) {
+      clamped[b] = fmaxf(r2[b], kSeriesSquared);
+    }
+    les_closed_form(clamped, value, slope);
+#pragma unroll
+    for (int b = 0; b < kBatch; ++b) {
+      if (r2[b] < kSeriesSquared) {
+        value[b] = even_series(kValueSeries, r2[b]);
+        slope[b] = even_series(kSlopeSeries, r2[b]);
+      }
+    }
+  } else {
+#pragma unroll
+    for (int b = 0; b < kBatch; ++b) {
+      value[b] = 0.f;
+      slope[b] = 0.f;
+    }
+    for (int k = 0; k < gaussian_count; ++k) {
+      const float4 coefficient = gaussians[k];
+#pragma unroll
+      for (int b = 0; b < kBatch; ++b) {
+        const float weight = exp2f(coefficient.y * r2[b]);
+        value[b] = fmaf(coefficient.x, weight, value[b]);
+        slope[b] = fmaf(coefficient.z, weight, slope[b]);
+      }
+    }
+  }
+}
+
+// Partners a tile holds besides those of the frame of the node.
+enum Exclusion : int {
+  // Every partner belongs to the frame and differs from the node.
+  kNoExclusion = 0,
+  // The node itself is among the partners.
+  kExcludeSelf = 1,
+  // The partners may also lie in other frames.
+  kExcludeOutsideFrame = 2,
+};
+
+// The state of the node a thread owns.
+template 
+struct OwnNode {
+  int index;
+  int low;   // first node of its frame
+  int high;  // past the last node of its frame
+  float4 head;
+  float4 tail;
+  float charge[kChannels];
+};
+
+// Adds the unscaled potential, force direction and virial that the partners
+// of one staged tile give the node to ``sum``.
+template 
+__device__ __forceinline__ void accumulate_tile(
+    const float4* __restrict__ tile_head,
+    const float4* __restrict__ tile_tail,
+    const float* __restrict__ tile_charge,
+    const float4* __restrict__ gaussians,
+    int gaussian_count,
+    int start,
+    const OwnNode& own,
+    float (&sum)[sum_width()]) {
+#pragma unroll 4
+  for (int base = 0; base < kTile; base += kBatch) {
+    float dx[kBatch];
+    float dy[kBatch];
+    float dz[kBatch];
+    float r2[kBatch];
+#pragma unroll
+    for (int b = 0; b < kBatch; ++b) {
+      const float4 other_head = tile_head[base + b];
+      dx[b] = own.head.x - other_head.x;
+      dy[b] = own.head.y - other_head.y;
+      dz[b] = own.head.z - other_head.z;
+    }
+    float value[kBatch];
+    float slope[kBatch];
+    if (far_batch(dx, dy, dz, r2)) {
+#pragma unroll
+      for (int b = 0; b < kBatch; ++b) {
+        const float inverse = rsqrtf(r2[b]);
+        value[b] = inverse;
+        slope[b] = -inverse * (inverse * inverse);
+      }
+    } else {
+#pragma unroll
+      for (int b = 0; b < kBatch; ++b) {
+        const float4 other_tail = tile_tail[base + b];
+        dx[b] += own.tail.x - other_tail.x;
+        dy[b] += own.tail.y - other_tail.y;
+        dz[b] += own.tail.z - other_tail.z;
+        // Both kernels are finite at r2 = 0, which a node paired with itself
+        // and coincident padding nodes reach with a zero weight.
+        r2[b] = fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b]));
+      }
+      evaluate_kernel(r2, gaussians, gaussian_count, value, slope);
+    }
+#pragma unroll
+    for (int b = 0; b < kBatch; ++b) {
+      const int other = start + base + b;
+      float weight = 1.f;
+      if constexpr (kExclusion == kExcludeSelf) {
+        weight = other != own.index ? 1.f : 0.f;
+      } else if constexpr (kExclusion == kExcludeOutsideFrame) {
+        weight = (other != own.index &&
+                  static_cast(other - own.low) <
+                      static_cast(own.high - own.low))
+                     ? 1.f
+                     : 0.f;
+      }
+      float product = 0.f;
+#pragma unroll
+      for (int c = 0; c < kChannels; ++c) {
+        const float other_charge =
+            kExclusion == kNoExclusion
+                ? tile_charge[(base + b) * kChannels + c]
+                : weight * tile_charge[(base + b) * kChannels + c];
+        sum[c] = fmaf(other_charge, value[b], sum[c]);
+        product = fmaf(own.charge[c], other_charge, product);
+      }
+      const float strength = product * slope[b];
+      const float sx = strength * dx[b];
+      const float sy = strength * dy[b];
+      const float sz = strength * dz[b];
+      sum[kChannels + 0] += sx;
+      sum[kChannels + 1] += sy;
+      sum[kChannels + 2] += sz;
+      sum[kChannels + 3] = fmaf(sx, dx[b], sum[kChannels + 3]);
+      sum[kChannels + 4] = fmaf(sy, dy[b], sum[kChannels + 4]);
+      sum[kChannels + 5] = fmaf(sz, dz[b], sum[kChannels + 5]);
+      sum[kChannels + 6] = fmaf(sx, dy[b], sum[kChannels + 6]);
+      sum[kChannels + 7] = fmaf(sx, dz[b], sum[kChannels + 7]);
+      sum[kChannels + 8] = fmaf(sy, dz[b], sum[kChannels + 8]);
+    }
+  }
+}
+
+// Sums, for every owner node of one block and every partner of one chunk, the
+// unscaled potential, force direction and virial of the node, and writes them
+// to ``partial[chunk][component][node - owner_begin]``. The owners are the
+// nodes [owner_begin, owner_begin + owner_count). A chunk is the part
+// blockIdx.y of gridDim.y equal parts of the partner range of the block or,
+// with ``chunk_tiles`` > 0, its chunk_tiles tiles from tile blockIdx.y *
+// chunk_tiles of the range.
+template 
+__global__ void __launch_bounds__(kTile)
+    direct_sum_kernel(const float4* __restrict__ head,
+                      const float4* __restrict__ tail,
+                      const float* __restrict__ charge,
+                      const int* __restrict__ frame,
+                      const int64_t* __restrict__ frame_ptr,
+                      int owner_begin,
+                      int owner_count,
+                      const double* __restrict__ coefficient,
+                      int gaussian_count,
+                      int chunk_tiles,
+                      float* __restrict__ partial) {
+  constexpr int kWidth = sum_width();
+  __shared__ float4 tile_head[kTile];
+  __shared__ float4 tile_tail[kTile];
+  __shared__ float tile_charge[kTile * kChannels];
+  __shared__ float4 gaussians[kMaxGaussians];
+  __shared__ int range[2];
+  __shared__ bool single_frame;
+
+  const int owner_end = owner_begin + owner_count;
+  const int first = owner_begin + blockIdx.x * kTile;
+  const int node = first + threadIdx.x;
+  const int own_frame = node < owner_end ? frame[node] : -1;
+
+  if constexpr (kKernel == kSog) {
+    for (int k = threadIdx.x; k < gaussian_count; k += kTile) {
+      const float amplitude = static_cast(coefficient[k]);
+      const float rate = static_cast(coefficient[gaussian_count + k]);
+      gaussians[k] = make_float4(amplitude, -rate,
+                                 -1.38629436111989061f * rate * amplitude, 0.f);
+    }
+  }
+  // The partners of the block are the nodes of the frames its nodes belong
+  // to: a contiguous range, since frames are contiguous on the node axis and
+  // padding nodes trail them.
+  if (threadIdx.x == 0) {
+    int last = min(first + kTile, owner_end) - 1;
+    while (last > first && frame[last] < 0) {
+      --last;
+    }
+    const int low_frame = frame[first];
+    const int high_frame = frame[last];
+    range[0] = low_frame < 0 ? 0 : static_cast(frame_ptr[low_frame]);
+    range[1] = low_frame < 0 ? 0 : static_cast(frame_ptr[high_frame + 1]);
+    // A block whose nodes all lie in one frame finds every partner of its
+    // range in that frame, which spares the frame test of every pair.
+    single_frame = low_frame >= 0 && low_frame == high_frame &&
+                   min(first + kTile, owner_end) - 1 == last;
+  }
+  __syncthreads();
+
+  // The chunk of the partner range this block sums, in whole tiles.
+  const int64_t tile_count = (range[1] - range[0] + kTile - 1) / kTile;
+  int64_t first_tile = tile_count * blockIdx.y / gridDim.y;
+  int64_t end_tile = tile_count * (blockIdx.y + 1) / gridDim.y;
+  if (chunk_tiles > 0) {
+    first_tile = min(tile_count, int64_t{blockIdx.y} * chunk_tiles);
+    end_tile = min(tile_count, first_tile + chunk_tiles);
+  }
+  const int begin = range[0] + static_cast(first_tile) * kTile;
+  const int end = min(range[1], range[0] + static_cast(end_tile) * kTile);
+
+  OwnNode own;
+  own.index = node;
+  own.low = 0;
+  own.high = 0;
+  own.head = make_float4(0.f, 0.f, 0.f, 0.f);
+  own.tail = own.head;
+  if (own_frame >= 0) {
+    own.head = head[node];
+    own.tail = tail[node];
+    own.low = static_cast(frame_ptr[own_frame]);
+    own.high = static_cast(frame_ptr[own_frame + 1]);
+  }
+#pragma unroll
+  for (int c = 0; c < kChannels; ++c) {
+    own.charge[c] = own_frame >= 0 ? charge[node * kChannels + c] : 0.f;
+  }
+
+  CompensatedSum total[kWidth];
+  for (int start = begin; start < end; start += kTile) {
+    const int partner = start + threadIdx.x;
+    const bool staged = partner < end;
+    // A slot past the chunk carries a zero charge and contributes nothing.
+    tile_head[threadIdx.x] =
+        staged ? head[partner] : make_float4(0.f, 0.f, 0.f, 0.f);
+    tile_tail[threadIdx.x] =
+        staged ? tail[partner] : make_float4(0.f, 0.f, 0.f, 0.f);
+#pragma unroll
+    for (int c = 0; c < kChannels; ++c) {
+      tile_charge[threadIdx.x * kChannels + c] =
+          staged ? charge[partner * kChannels + c] : 0.f;
+    }
+    __syncthreads();
+
+    float sum[kWidth];
+#pragma unroll
+    for (int k = 0; k < kWidth; ++k) {
+      sum[k] = 0.f;
+    }
+    // In a single-frame block only the tiles that overlap the block's own
+    // nodes hold a node itself, so the tests are confined to them and to the
+    // blocks that span frames; the choice is uniform over the block.
+    if (!single_frame) {
+      accumulate_tile(
+          tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start,
+          own, sum);
+    } else if (start < first + kTile && first < start + kTile) {
+      accumulate_tile(
+          tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start,
+          own, sum);
+    } else {
+      accumulate_tile(
+          tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start,
+          own, sum);
+    }
+#pragma unroll
+    for (int k = 0; k < kWidth; ++k) {
+      total[k].add(sum[k]);
+    }
+    __syncthreads();
+  }
+
+  if (node < owner_end) {
+    float* out = partial +
+                 static_cast(blockIdx.y) * kWidth * owner_count +
+                 (node - owner_begin);
+#pragma unroll
+    for (int k = 0; k < kWidth; ++k) {
+      out[static_cast(k) * owner_count] = total[k].sum;
+    }
+  }
+}
+
+// Adds the chunk sums of one component of one owner node in chunk order, one
+// thread per component and owner, and scales the result into the potential,
+// the force or the per-node virial of the owner. With positions scaled by
+// alpha, K(r) = alpha f(x), K'(r) / r = alpha^3 f'(x) / x and d = d' / alpha,
+// so the potential carries alpha, the force alpha^2 and the virial alpha.
+template 
+__global__ void reduce_chunks(const float* __restrict__ partial,
+                              int owner_count,
+                              int chunk_count,
+                              double prefactor,
+                              const double* __restrict__ alpha,
+                              float* __restrict__ potential,
+                              float* __restrict__ force,
+                              float* __restrict__ atom_virial) {
+  constexpr int kWidth = sum_width();
+  const int64_t index =
+      static_cast(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= static_cast(kWidth) * owner_count) {
+    return;
+  }
+  const int component = static_cast(index / owner_count);
+  const int node = static_cast(index - component * int64_t(owner_count));
+  double value = 0.0;
+  for (int chunk = 0; chunk < chunk_count; ++chunk) {
+    value +=
+        partial[static_cast(chunk) * kWidth * owner_count + index];
+  }
+  const double unit = alpha == nullptr ? 1.0 : *alpha;
+  if (component < kChannels) {
+    potential[node * kChannels + component] =
+        static_cast(prefactor * unit * value);
+  } else if (component < kChannels + 3) {
+    force[node * 3 + component - kChannels] =
+        static_cast(-prefactor * unit * unit * value);
+  } else {
+    // Each pair term is split half to each of its two nodes. The six sums
+    // are xx, yy, zz, xy, xz, yz.
+    constexpr int kRow[6] = {0, 1, 2, 0, 0, 1};
+    constexpr int kColumn[6] = {0, 1, 2, 1, 2, 2};
+    const int pair = component - kChannels - 3;
+    const float entry = static_cast(-0.5 * prefactor * unit * value);
+    float* out = atom_virial + static_cast(node) * 9;
+    out[3 * kRow[pair] + kColumn[pair]] = entry;
+    out[3 * kColumn[pair] + kRow[pair]] = entry;
+  }
+}
+
+// Sums, over one slice of one frame, the per-node virial and, with a
+// constraint, the potential of the charged nodes in the constrained channel.
+template 
+__global__ void __launch_bounds__(kFrameThreads)
+    virial_slices(const float* __restrict__ atom_virial,
+                  const float* __restrict__ potential,
+                  const bool* __restrict__ node_mask,
+                  const int64_t* __restrict__ frame_ptr,
+                  int channels,
+                  int channel,
+                  double* __restrict__ slice_sum) {
+  constexpr int kComponents = kConstrained ? 10 : 9;
+  const NodeRange slice = slice_of_frame(frame_ptr);
+  double value[kComponents] = {};
+  for (int64_t node = slice.begin + threadIdx.x; node < slice.end;
+       node += kFrameThreads) {
+#pragma unroll
+    for (int k = 0; k < 9; ++k) {
+      value[k] += atom_virial[node * 9 + k];
+    }
+    if constexpr (kConstrained) {
+      if (node_mask[node]) {
+        value[9] += potential[node * channels + channel];
+      }
+    }
+  }
+  write_slice_sum(value, slice_sum);
+}
+
+// Adds the slice sums of virial_slices into the virial of every frame and,
+// with a constraint, the potential summed over its charged nodes.
+template 
+__global__ void virial_total(const double* __restrict__ slice_sum,
+                             int slice_count,
+                             float* __restrict__ virial,
+                             double* __restrict__ frame_stat) {
+  constexpr int kComponents = kConstrained ? 10 : 9;
+  if (threadIdx.x >= kComponents) {
+    return;
+  }
+  double value = 0.0;
+  for (int slice = 0; slice < slice_count; ++slice) {
+    value +=
+        slice_sum[(static_cast(blockIdx.x) * slice_count + slice) *
+                      kComponents +
+                  threadIdx.x];
+  }
+  if (threadIdx.x < 9) {
+    virial[blockIdx.x * 9 + threadIdx.x] = static_cast(value);
+  } else {
+    frame_stat[blockIdx.x * kStatWidth + kStatPotential] = value;
+  }
+}
+
+// Writes the atomic energy q_i . phi_i / 2 of every owner node and the
+// derivative of the energy with respect to its raw charges: w_i (phi_i -
+// mean_f) in the constrained channel and w_i phi_i in the others, w_i marking
+// the charged nodes. ``pair_charge``, ``node_mask`` and ``frame`` start at the
+// first owner.
+__global__ void project_charges(const float* __restrict__ potential,
+                                const float* __restrict__ pair_charge,
+                                const bool* __restrict__ node_mask,
+                                const int* __restrict__ frame,
+                                const double* __restrict__ frame_stat,
+                                int owner_count,
+                                int channels,
+                                int channel,
+                                double* __restrict__ atom_energy,
+                                float* __restrict__ charge_gradient) {
+  const int node = blockIdx.x * blockDim.x + threadIdx.x;
+  if (node >= owner_count) {
+    return;
+  }
+  const int f = frame[node];
+  const bool charged = f >= 0 && node_mask[node];
+  const double mean = charged && frame_stat != nullptr
+                          ? frame_stat[f * kStatWidth + kStatPotential] /
+                                frame_stat[f * kStatWidth + kStatCount]
+                          : 0.0;
+  double energy = 0.0;
+  for (int c = 0; c < channels; ++c) {
+    const int64_t index = static_cast(node) * channels + c;
+    const double phi = potential[index];
+    energy += static_cast(pair_charge[index]) * phi;
+    const double gradient = charged ? phi : 0.0;
+    charge_gradient[index] =
+        static_cast(c == channel ? gradient - mean : gradient);
+  }
+  atom_energy[node] = 0.5 * energy;
+}
+
+// The staged partners, one row per node, and the sums of the owners.
+struct DirectSumBuffers {
+  const float4* head;
+  const float4* tail;
+  const float* charge;
+  const int* frame;
+  const int64_t* frame_ptr;
+  int node_count;
+  const double* coefficient;
+  int gaussian_count;
+  float* potential;
+  float* force;
+  float* atom_virial;
+};
+
+// The launch of the pair kernel: whether it takes narrow tiles, and its
+// chunks of the partner axis, ``chunk_count`` of them and, for narrow tiles,
+// ``chunk_tiles`` tiles each (direct_sum_kernel).
+struct PairLaunch {
+  bool narrow;
+  int chunk_count;
+  int chunk_tiles;
+};
+
+// The launch for the given counts of owners and partners. The blocks of one
+// launch of wide tiles all carry the same work, so the launch fills the
+// resident block slots of the device several times over, which bounds the
+// idle share of the last wave. A frame that spans every partner has at most
+// one chunk per kWideChunkTiles of its tiles, rounded up, so that in a launch
+// with slots to spare no chunk holds more than kWideChunkTiles tiles: the
+// longest block sets the duration of such a launch. A system whose wide
+// blocks would fill less than a third of the resident block slots leaves too
+// few warps on the multiprocessors to hide the latency of the pair loop, and
+// takes narrow tiles in chunks of kNarrowChunkTiles, as many as a frame that
+// spans every partner holds.
+PairLaunch pair_launch(int owner_count,
+                       int partner_count,
+                       int resident_blocks_per_multiprocessor) {
+  const int multiprocessors =
+      at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  const int owner_blocks = (owner_count + kWideTile - 1) / kWideTile;
+  const int partner_tiles = (partner_count + kWideTile - 1) / kWideTile;
+  const int64_t wanted_blocks = static_cast(kWaves) *
+                                resident_blocks_per_multiprocessor *
+                                multiprocessors;
+  const int wanted =
+      static_cast((wanted_blocks + owner_blocks - 1) / owner_blocks);
+  const int chunk_limit =
+      (partner_tiles + kWideChunkTiles - 1) / kWideChunkTiles;
+  const int chunk_count = std::max(1, std::min(wanted, chunk_limit));
+  if (3 * int64_t{owner_blocks} * chunk_count >=
+      int64_t{resident_blocks_per_multiprocessor} * multiprocessors) {
+    return {false, chunk_count, 0};
+  }
+  const int narrow_tiles = (partner_count + kNarrowTile - 1) / kNarrowTile;
+  return {
+      true,
+      std::max(1, (narrow_tiles + kNarrowChunkTiles - 1) / kNarrowChunkTiles),
+      kNarrowChunkTiles};
+}
+
+// The blocks of wide tiles that a multiprocessor of the current device keeps
+// resident. The occupancy is a property of the compiled kernel and of the
+// device, so each device of the process is queried once.
+template 
+int resident_blocks_per_multiprocessor() {
+  static std::mutex mutex;
+  static std::map blocks;
+  const int device = c10::cuda::current_device();
+  const std::lock_guard lock(mutex);
+  const auto known = blocks.find(device);
+  if (known != blocks.end()) {
+    return known->second;
+  }
+  int value = 0;
+  C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+      &value, direct_sum_kernel, kWideTile, 0));
+  return blocks.emplace(device, std::max(1, value)).first->second;
+}
+
+template 
+void launch_tiles(const DirectSumBuffers& buffers,
+                  int owner_begin,
+                  int owner_count,
+                  const PairLaunch& pair,
+                  double prefactor,
+                  cudaStream_t stream) {
+  constexpr int kWidth = sum_width();
+  auto partial = torch::empty(
+      {static_cast(pair.chunk_count) * owner_count * kWidth},
+      torch::TensorOptions()
+          .dtype(torch::kFloat32)
+          .device(torch::kCUDA, c10::cuda::current_device()));
+  const dim3 grid((owner_count + kTile - 1) / kTile, pair.chunk_count);
+  direct_sum_kernel<<>>(
+      buffers.head, buffers.tail, buffers.charge, buffers.frame,
+      buffers.frame_ptr, owner_begin, owner_count, buffers.coefficient,
+      buffers.gaussian_count, pair.chunk_tiles, partial.data_ptr());
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  const int64_t reduce_threads = static_cast(kWidth) * owner_count;
+  reduce_chunks
+      <<((reduce_threads + 255) / 256), 256, 0, stream>>>(
+          partial.data_ptr(), owner_count, pair.chunk_count, prefactor,
+          kKernel == kLes ? buffers.coefficient : nullptr, buffers.potential,
+          buffers.force, buffers.atom_virial);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+template 
+void launch(const DirectSumBuffers& buffers,
+            int owner_begin,
+            int owner_count,
+            double prefactor,
+            cudaStream_t stream) {
+  const PairLaunch pair =
+      pair_launch(owner_count, buffers.node_count,
+                  resident_blocks_per_multiprocessor());
+  if (pair.narrow) {
+    launch_tiles(
+        buffers, owner_begin, owner_count, pair, prefactor, stream);
+  } else {
+    launch_tiles(
+        buffers, owner_begin, owner_count, pair, prefactor, stream);
+  }
+}
+
+static_assert(kMaxChannels == 8,
+              "dispatch_channels instantiates the channel counts 1 to 8");
+
+template 
+void dispatch_channels(int channels,
+                       const DirectSumBuffers& buffers,
+                       int owner_begin,
+                       int owner_count,
+                       double prefactor,
+                       cudaStream_t stream) {
+  switch (channels) {
+    case 1:
+      return launch(buffers, owner_begin, owner_count, prefactor,
+                                stream);
+    case 2:
+      return launch(buffers, owner_begin, owner_count, prefactor,
+                                stream);
+    case 3:
+      return launch(buffers, owner_begin, owner_count, prefactor,
+                                stream);
+    case 4:
+      return launch(buffers, owner_begin, owner_count, prefactor,
+                                stream);
+    case 5:
+      return launch(buffers, owner_begin, owner_count, prefactor,
+                                stream);
+    case 6:
+      return launch(buffers, owner_begin, owner_count, prefactor,
+                                stream);
+    case 7:
+      return launch(buffers, owner_begin, owner_count, prefactor,
+                                stream);
+    case 8:
+      return launch(buffers, owner_begin, owner_count, prefactor,
+                                stream);
+    default:
+      TORCH_CHECK(false, "latent_charge_sum: unsupported channel count ",
+                  channels);
+  }
+}
+
+}  // namespace
+
+namespace deepmd {
+namespace long_range {
+
+// Evaluates the owner nodes [owner_begin, owner_begin + owner_count) against
+// the nodes of their frames. Returns the outputs of the owners, row r for
+// owner owner_begin + r, and the per-frame virial of the owners. Owners other
+// than every node lie in a single frame. A non-null ``communicator`` spans the
+// domains whose owners share that frame, over which the potential of the
+// constraint is summed.
+ChargeSumOutputs direct_sum(const torch::Tensor& coord,
+                            const torch::Tensor& charge,
+                            const torch::Tensor& node_mask,
+                            const torch::Tensor& n_node_per_frame,
+                            const torch::Tensor& total_charge,
+                            const torch::Tensor& coefficient,
+                            int64_t kernel,
+                            int64_t channel,
+                            double prefactor,
+                            int owner_begin,
+                            int owner_count,
+                            const torch::Tensor* communicator) {
+  TORCH_CHECK(coord.size(0) < std::numeric_limits::max() / kThreads,
+              "latent_charge_sum: too many nodes");
+  const int node_count = static_cast(coord.size(0));
+  const int frame_count = static_cast(n_node_per_frame.size(0));
+  const int64_t channels = charge.size(1);
+  const auto float_options = charge.options();
+  const auto double_options = coord.options();
+  const bool every_node = owner_begin == 0 && owner_count == node_count;
+  const bool constrained = channel >= 0;
+  const int constrained_channel = static_cast(channel);
+  // The kernels write every entry of the outputs.
+  auto atom_energy = torch::empty({owner_count, 1}, double_options);
+  auto charge_gradient = torch::empty({owner_count, channels}, float_options);
+  auto force = torch::empty({owner_count, 3}, float_options);
+  auto atom_virial = torch::empty({owner_count, 9}, float_options);
+  auto virial = torch::empty({frame_count, 9}, float_options);
+
+  const auto counts = n_node_per_frame.to(torch::kInt64).contiguous();
+  const int slice_count = frame_slices(node_count, frame_count);
+  const int owner_slice_count = frame_slices(owner_count, frame_count);
+  const int components = constrained ? 10 : 9;
+  Workspace workspace;
+  const int64_t frame_ptr_at = workspace.reserve(frame_count + 1);
+  const int64_t head_at = workspace.reserve(node_count);
+  const int64_t tail_at = workspace.reserve(node_count);
+  const int64_t frame_at = workspace.reserve(node_count);
+  const int64_t pair_charge_at =
+      workspace.reserve(int64_t{node_count} * channels);
+  const int64_t potential_at =
+      workspace.reserve(int64_t{owner_count} * channels);
+  const int64_t frame_stat_at =
+      workspace.reserve(int64_t{frame_count} * kStatWidth);
+  const int64_t slice_charge_at =
+      workspace.reserve(int64_t{frame_count} * slice_count * 2);
+  const int64_t slice_virial_at = workspace.reserve(
+      int64_t{frame_count} * owner_slice_count * components);
+  workspace.allocate(coord.device());
+  auto* frame_ptr = workspace.at(frame_ptr_at);
+  auto* head = workspace.at(head_at);
+  auto* tail = workspace.at(tail_at);
+  auto* frame = workspace.at(frame_at);
+  auto* pair_charge = workspace.at(pair_charge_at);
+  auto* potential = workspace.at(potential_at);
+  auto* frame_stat =
+      constrained ? workspace.at(frame_stat_at) : nullptr;
+  auto* slice_charge = workspace.at(slice_charge_at);
+  auto* slice_virial = workspace.at(slice_virial_at);
+
+  const auto stream = at::cuda::getCurrentCUDAStream();
+  frame_offsets<<<1, kOffsetThreads, 0, stream>>>(counts.data_ptr(),
+                                                  frame_count, frame_ptr);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  if (constrained) {
+    charge_slices<<>>(
+        charge.data_ptr(), node_mask.data_ptr(), frame_ptr,
+        static_cast(channels), constrained_channel, slice_charge);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads,
+                   kFrameThreads, 0, stream>>>(
+        slice_charge, slice_count, frame_count, total_charge.data_ptr(),
+        frame_stat);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  }
+  // Alpha stays on the device, so that a call never waits for it.
+  const double* alpha =
+      kernel == kLes ? coefficient.data_ptr() : nullptr;
+  prepare_nodes<<<(node_count + kThreads - 1) / kThreads, kThreads, 0,
+                  stream>>>(coord.data_ptr(), charge.data_ptr(),
+                            node_mask.data_ptr(), frame_ptr, frame_count,
+                            node_count, static_cast(channels),
+                            constrained_channel, alpha, frame_stat, head, tail,
+                            frame, pair_charge);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+  if (owner_count > 0) {
+    const DirectSumBuffers buffers{head,
+                                   tail,
+                                   pair_charge,
+                                   frame,
+                                   frame_ptr,
+                                   node_count,
+                                   coefficient.data_ptr(),
+                                   static_cast(coefficient.numel() / 2),
+                                   potential,
+                                   force.data_ptr(),
+                                   atom_virial.data_ptr()};
+    if (kernel == kLes) {
+      dispatch_channels(static_cast(channels), buffers, owner_begin,
+                              owner_count, prefactor, stream.stream());
+    } else {
+      dispatch_channels(static_cast(channels), buffers, owner_begin,
+                              owner_count, prefactor, stream.stream());
+    }
+  }
+  torch::Tensor owner_frame_storage;
+  const int64_t* owner_frame_ptr =
+      owner_frame_offsets(every_node, frame_ptr, owner_count, counts.options(),
+                          owner_frame_storage);
+  const bool* owner_mask = node_mask.data_ptr() + owner_begin;
+  const dim3 owner_slice_grid(frame_count, owner_slice_count);
+  if (constrained) {
+    virial_slices<<>>(
+        atom_virial.data_ptr(), potential, owner_mask, owner_frame_ptr,
+        static_cast(channels), constrained_channel, slice_virial);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    virial_total<<>>(
+        slice_virial, owner_slice_count, virial.data_ptr(), frame_stat);
+  } else {
+    virial_slices<<>>(
+        atom_virial.data_ptr(), potential, owner_mask, owner_frame_ptr,
+        static_cast(channels), constrained_channel, slice_virial);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    virial_total<<>>(
+        slice_virial, owner_slice_count, virial.data_ptr(), nullptr);
+  }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  if (constrained && communicator != nullptr) {
+    // The mean potential of the constraint runs over the charged nodes of
+    // every domain.
+    auto stat =
+        torch::from_blob(frame_stat, {frame_count, kStatWidth}, double_options);
+    auto potential_sum = stat.select(1, kStatPotential);
+    potential_sum.copy_(
+        sum_over_domains(*communicator, potential_sum.contiguous()));
+  }
+  if (owner_count > 0) {
+    project_charges<<<(owner_count + kThreads - 1) / kThreads, kThreads, 0,
+                      stream>>>(
+        potential, pair_charge + int64_t{owner_begin} * channels, owner_mask,
+        frame + owner_begin, frame_stat, owner_count,
+        static_cast(channels), constrained_channel,
+        atom_energy.data_ptr(), charge_gradient.data_ptr());
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  }
+  return {atom_energy, charge_gradient, force, atom_virial, virial};
+}
+
+}  // namespace long_range
+}  // namespace deepmd
diff --git a/source/op/pt/long_range/domain_exchange.cc b/source/op/pt/long_range/domain_exchange.cc
new file mode 100644
index 0000000000..cbe4f07a23
--- /dev/null
+++ b/source/op/pt/long_range/domain_exchange.cc
@@ -0,0 +1,165 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// MPI exchanges of the latent-charge long range (contract in
+// domain_exchange.h).
+//
+// With an MPI library that reads device memory (MPIX_Query_cuda_support), the
+// rows travel between device buffers once the current stream has finished
+// writing them; otherwise they are staged through pinned host buffers. Either
+// way the host waits for the device once per exchange, before MPI reads the
+// data; the results return to the device without a further wait.
+
+#include "domain_exchange.h"
+
+#ifdef USE_MPI
+#include 
+#ifdef OMPI_MPI_H
+#include 
+#endif
+#endif
+#include 
+#include 
+
+#include 
+#include 
+
+namespace deepmd {
+namespace long_range {
+namespace {
+
+std::int64_t address_of(const torch::Tensor& communicator) {
+  TORCH_CHECK(communicator.device().is_cpu() &&
+                  communicator.scalar_type() == torch::kInt64 &&
+                  communicator.numel() == 1,
+              "latent_charge_sum: communicator must be an int64 CPU "
+              "tensor with one element");
+  return communicator.contiguous().data_ptr()[0];
+}
+
+#ifdef USE_MPI
+// The handle of the caller, of the width of MPI_Comm in this MPI library: an
+// int in MPICH, a pointer in Open MPI.
+MPI_Comm communicator_of(const torch::Tensor& communicator) {
+  static_assert(sizeof(std::uintptr_t) <= sizeof(std::int64_t),
+                "an address must fit in an int64 tensor");
+  return *reinterpret_cast(
+      static_cast(address_of(communicator)));
+}
+
+bool reads_device_memory() {
+#ifdef NO_CUDA_AWARE
+  return false;
+#else
+  return MPIX_Query_cuda_support() != 0;
+#endif
+}
+#endif
+
+}  // namespace
+
+int domain_count(const torch::Tensor& communicator) {
+  if (address_of(communicator) == 0) {
+    return 1;
+  }
+#ifdef USE_MPI
+  int initialized = 0;
+  MPI_Initialized(&initialized);
+  TORCH_CHECK(initialized != 0,
+              "latent_charge_sum: a communicator was given, but MPI is "
+              "not initialized");
+  int size = 1;
+  MPI_Comm_size(communicator_of(communicator), &size);
+  return size;
+#else
+  TORCH_CHECK(false,
+              "latent_charge_sum: an evaluation across MPI ranks "
+              "requires DeePMD-kit built with MPI");
+#endif
+}
+
+std::pair gather_domain_rows(
+    const torch::Tensor& communicator, const torch::Tensor& rows) {
+#ifdef USE_MPI
+  TORCH_CHECK(rows.is_cuda() && rows.dim() == 2 &&
+                  rows.scalar_type() == torch::kFloat64 && rows.is_contiguous(),
+              "gather_domain_rows: rows must be a contiguous float64 CUDA "
+              "matrix");
+  const MPI_Comm comm = communicator_of(communicator);
+  int rank = 0;
+  int size = 1;
+  MPI_Comm_rank(comm, &rank);
+  MPI_Comm_size(comm, &size);
+  const std::int64_t width = rows.size(1);
+  std::int64_t local = rows.size(0);
+  std::vector counts(size);
+  MPI_Allgather(&local, 1, MPI_INT64_T, counts.data(), 1, MPI_INT64_T, comm);
+  std::vector receive_counts(size);
+  std::vector displacements(size);
+  std::int64_t total = 0;
+  std::int64_t offset = 0;
+  for (int r = 0; r < size; ++r) {
+    if (r == rank) {
+      offset = total;
+    }
+    TORCH_CHECK((total + counts[r]) * width <= std::numeric_limits::max(),
+                "gather_domain_rows: the rows of all ranks exceed the int "
+                "range of an MPI count");
+    receive_counts[r] = static_cast(counts[r] * width);
+    displacements[r] = static_cast(total * width);
+    total += counts[r];
+  }
+  const int send_count = static_cast(local * width);
+  auto gathered = torch::empty({total, width}, rows.options());
+  if (reads_device_memory()) {
+    at::cuda::getCurrentCUDAStream().synchronize();
+    MPI_Allgatherv(rows.data_ptr(), send_count, MPI_DOUBLE,
+                   gathered.data_ptr(), receive_counts.data(),
+                   displacements.data(), MPI_DOUBLE, comm);
+  } else {
+    const auto host = rows.options().device(torch::kCPU).pinned_memory(true);
+    auto host_rows = torch::empty({local, width}, host);
+    host_rows.copy_(rows);
+    auto host_gathered = torch::empty({total, width}, host);
+    MPI_Allgatherv(host_rows.data_ptr(), send_count, MPI_DOUBLE,
+                   host_gathered.data_ptr(), receive_counts.data(),
+                   displacements.data(), MPI_DOUBLE, comm);
+    gathered.copy_(host_gathered, /*non_blocking=*/true);
+  }
+  return {gathered, offset};
+#else
+  TORCH_CHECK(false,
+              "gather_domain_rows: an evaluation across MPI ranks requires "
+              "DeePMD-kit built with MPI");
+#endif
+}
+
+torch::Tensor sum_over_domains(const torch::Tensor& communicator,
+                               const torch::Tensor& values) {
+#ifdef USE_MPI
+  TORCH_CHECK(values.dim() == 1 && values.scalar_type() == torch::kFloat64,
+              "sum_over_domains: values must be a float64 vector");
+  const MPI_Comm comm = communicator_of(communicator);
+  int size = 1;
+  MPI_Comm_size(comm, &size);
+  const torch::Tensor local = values.to(torch::kCPU).contiguous();
+  const int count = static_cast(local.numel());
+  std::vector every(static_cast(size) * count);
+  MPI_Allgather(local.data_ptr(), count, MPI_DOUBLE, every.data(),
+                count, MPI_DOUBLE, comm);
+  auto total = torch::zeros({count}, local.options().pinned_memory(true));
+  double* sum = total.data_ptr();
+  for (int r = 0; r < size; ++r) {
+    for (int k = 0; k < count; ++k) {
+      sum[k] += every[static_cast(r) * count + k];
+    }
+  }
+  return total.to(values.device(), /*non_blocking=*/true);
+#else
+  TORCH_CHECK(false,
+              "sum_over_domains: an evaluation across MPI ranks requires "
+              "DeePMD-kit built with MPI");
+#endif
+}
+
+}  // namespace long_range
+}  // namespace deepmd
diff --git a/source/op/pt/long_range/domain_exchange.h b/source/op/pt/long_range/domain_exchange.h
new file mode 100644
index 0000000000..0ebd9698d4
--- /dev/null
+++ b/source/op/pt/long_range/domain_exchange.h
@@ -0,0 +1,38 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// Exchange of the latent-charge long range between the domains of a
+// decomposed system.
+//
+// A domain is one MPI rank of the communicator that a caller hands to
+// deepmd::latent_charge_sum as an int64 tensor holding the address of its
+// MPI_Comm handle. A zero address stands for a single domain. Both exchanges
+// follow the rank order, so every rank receives the same rows and the same
+// sums, independent of the reduction algorithm of the MPI library.
+
+#pragma once
+
+#include 
+
+#include 
+#include 
+
+namespace deepmd {
+namespace long_range {
+
+// The number of ranks of ``communicator``; one for a zero address.
+int domain_count(const torch::Tensor& communicator);
+
+// Concatenates, in rank order, the rows of every rank. ``rows`` is a
+// contiguous float64 matrix of shape (M_r, W) on the CUDA device, with the
+// same W on every rank. Returns the rows of all ranks, of shape (sum_r M_r,
+// W) on the device of ``rows``, and the index of the first row of this rank.
+std::pair gather_domain_rows(
+    const torch::Tensor& communicator, const torch::Tensor& rows);
+
+// Sums the float64 vector ``values`` over the ranks, adding the vectors of
+// the ranks in rank order. Returns the sum on the device of ``values``.
+torch::Tensor sum_over_domains(const torch::Tensor& communicator,
+                               const torch::Tensor& values);
+
+}  // namespace long_range
+}  // namespace deepmd
diff --git a/source/op/pt/long_range/frame_sums.cuh b/source/op/pt/long_range/frame_sums.cuh
new file mode 100644
index 0000000000..967277d1bc
--- /dev/null
+++ b/source/op/pt/long_range/frame_sums.cuh
@@ -0,0 +1,240 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// Per-frame bookkeeping shared by the CUDA units of the latent-charge long
+// range: the workspace of a call, the offsets of the frames on the node axis,
+// the sums over the nodes of every frame in a fixed order, and the statistics
+// of the charge constraint.
+//
+// The kernels are defined in an unnamed namespace, so that every unit that
+// includes this header compiles its own copy.
+
+#pragma once
+
+#include 
+#include 
+
+#include 
+#include 
+
+namespace {
+
+// The intermediate buffers of one call, carved out of a single device
+// allocation: ``reserve`` lays out a buffer of ``count`` elements and returns
+// its byte offset, and ``at`` its address once ``allocate`` has made the
+// allocation. Every buffer starts on a 256-byte boundary.
+class Workspace {
+ public:
+  template 
+  int64_t reserve(int64_t count) {
+    bytes_ = (bytes_ + kAlignment - 1) / kAlignment * kAlignment;
+    const int64_t offset = bytes_;
+    bytes_ += count * static_cast(sizeof(T));
+    return offset;
+  }
+
+  void allocate(const torch::Device& device) {
+    storage_ = torch::empty(
+        {std::max(bytes_, 1)},
+        torch::TensorOptions().dtype(torch::kUInt8).device(device));
+  }
+
+  template 
+  T* at(int64_t offset) const {
+    return reinterpret_cast(storage_.data_ptr() + offset);
+  }
+
+ private:
+  static constexpr int64_t kAlignment = 256;
+  int64_t bytes_ = 0;
+  torch::Tensor storage_;
+};
+
+// Index of the frame holding ``node``, or -1 for a padding node past every
+// frame. ``frame_ptr`` holds ``frame_count + 1`` offsets.
+__device__ __forceinline__ int frame_of(const int64_t* frame_ptr,
+                                        int frame_count,
+                                        int64_t node) {
+  if (node >= frame_ptr[frame_count]) {
+    return -1;
+  }
+  int low = 0;
+  int high = frame_count;
+  while (high - low > 1) {
+    const int middle = (low + high) / 2;
+    if (frame_ptr[middle] <= node) {
+      low = middle;
+    } else {
+      high = middle;
+    }
+  }
+  return low;
+}
+
+// Writes the offsets of the frames on the node axis: frame_ptr[f] counts the
+// nodes of the frames before f, for f in [0, frame_count]. One block scans
+// the counts, each thread a contiguous run of frames.
+constexpr int kOffsetThreads = 256;
+
+__global__ void __launch_bounds__(kOffsetThreads)
+    frame_offsets(const int64_t* __restrict__ n_node_per_frame,
+                  int frame_count,
+                  int64_t* __restrict__ frame_ptr) {
+  __shared__ int64_t run_total[kOffsetThreads];
+  const int run = (frame_count + kOffsetThreads - 1) / kOffsetThreads;
+  const int begin = min(frame_count, static_cast(threadIdx.x) * run);
+  const int end = min(frame_count, begin + run);
+  int64_t total = 0;
+  for (int f = begin; f < end; ++f) {
+    total += n_node_per_frame[f];
+  }
+  run_total[threadIdx.x] = total;
+  __syncthreads();
+  // Inclusive scan of the run totals.
+  for (int stride = 1; stride < kOffsetThreads; stride *= 2) {
+    const int64_t addend =
+        threadIdx.x >= stride ? run_total[threadIdx.x - stride] : 0;
+    __syncthreads();
+    run_total[threadIdx.x] += addend;
+    __syncthreads();
+  }
+  int64_t offset = threadIdx.x == 0 ? 0 : run_total[threadIdx.x - 1];
+  for (int f = begin; f < end; ++f) {
+    frame_ptr[f] = offset;
+    offset += n_node_per_frame[f];
+  }
+  if (threadIdx.x == kOffsetThreads - 1) {
+    frame_ptr[frame_count] = run_total[kOffsetThreads - 1];
+  }
+}
+
+// Per-frame sums. The nodes of a frame are split into ``gridDim.y``
+// contiguous slices, one block of kFrameThreads each, and a finalizing kernel
+// adds the float64 sums of the slices in slice order.
+constexpr int kFrameThreads = 256;
+// Nodes per slice of a frame that a slice pass aims for.
+constexpr int kSliceNodes = 2048;
+constexpr int kMaxSlices = 256;
+
+// The slices of every frame of a pass over ``nodes`` nodes of ``frame_count``
+// frames, from their average count per frame.
+inline int frame_slices(int64_t nodes, int64_t frame_count) {
+  const int64_t average = (nodes + frame_count - 1) / frame_count;
+  return static_cast(std::min(
+      kMaxSlices,
+      std::max(1, (average + kSliceNodes - 1) / kSliceNodes)));
+}
+
+// The offsets of the frames of the owners on the axis of the owners: those of
+// the frames, ``frame_ptr``, when every node is an owner, else those of the
+// one frame that the owners share, {0, owner_count}, which ``storage``
+// receives on the device of ``options``.
+inline const int64_t* owner_frame_offsets(bool every_node,
+                                          const int64_t* frame_ptr,
+                                          int owner_count,
+                                          const torch::TensorOptions& options,
+                                          torch::Tensor& storage) {
+  if (every_node) {
+    return frame_ptr;
+  }
+  storage = torch::full({2}, int64_t{owner_count}, options);
+  storage.select(0, 0).zero_();
+  return storage.data_ptr();
+}
+
+struct NodeRange {
+  int64_t begin;
+  int64_t end;
+};
+
+// The nodes of the slice of the frame that this block sums.
+__device__ __forceinline__ NodeRange
+slice_of_frame(const int64_t* __restrict__ frame_ptr) {
+  const int64_t frame_begin = frame_ptr[blockIdx.x];
+  const int64_t frame_end = frame_ptr[blockIdx.x + 1];
+  const int64_t length = (frame_end - frame_begin + gridDim.y - 1) / gridDim.y;
+  const int64_t begin = frame_begin + blockIdx.y * length;
+  return {begin, min(frame_end, begin + length)};
+}
+
+// Adds ``value`` over the threads of the block in a fixed order and writes the
+// totals to the entry of this slice in ``slice_sum``.
+template 
+__device__ __forceinline__ void write_slice_sum(
+    const double (&value)[kComponents], double* __restrict__ slice_sum) {
+  __shared__ double partial[kComponents][kFrameThreads];
+#pragma unroll
+  for (int k = 0; k < kComponents; ++k) {
+    partial[k][threadIdx.x] = value[k];
+  }
+  __syncthreads();
+  for (int stride = kFrameThreads / 2; stride > 0; stride /= 2) {
+    if (threadIdx.x < stride) {
+#pragma unroll
+      for (int k = 0; k < kComponents; ++k) {
+        partial[k][threadIdx.x] += partial[k][threadIdx.x + stride];
+      }
+    }
+    __syncthreads();
+  }
+  if (threadIdx.x < kComponents) {
+    slice_sum[(static_cast(blockIdx.x) * gridDim.y + blockIdx.y) *
+                  kComponents +
+              threadIdx.x] = partial[threadIdx.x][0];
+  }
+}
+
+// Per-frame statistics of the charge constraint: the shift of the
+// constrained channel, the number of charged nodes, and the potential of the
+// charged nodes in that channel summed over them.
+constexpr int kStatShift = 0;
+constexpr int kStatCount = 1;
+constexpr int kStatPotential = 2;
+constexpr int kStatWidth = 3;
+
+// Sums, over one slice of one frame, the number of charged nodes and their raw
+// charges in the constrained channel.
+__global__ void __launch_bounds__(kFrameThreads)
+    charge_slices(const float* __restrict__ charge,
+                  const bool* __restrict__ node_mask,
+                  const int64_t* __restrict__ frame_ptr,
+                  int channels,
+                  int channel,
+                  double* __restrict__ slice_sum) {
+  const NodeRange slice = slice_of_frame(frame_ptr);
+  double value[2] = {0.0, 0.0};
+  for (int64_t node = slice.begin + threadIdx.x; node < slice.end;
+       node += kFrameThreads) {
+    if (node_mask[node]) {
+      value[0] += 1.0;
+      value[1] += charge[node * channels + channel];
+    }
+  }
+  write_slice_sum<2>(value, slice_sum);
+}
+
+// The shift s_f = (sum_f c - Q_f) / max(n_f, 1) of the constrained channel of
+// every frame, from the slice sums of charge_slices.
+__global__ void charge_shift(const double* __restrict__ slice_sum,
+                             int slice_count,
+                             int frame_count,
+                             const double* __restrict__ total_charge,
+                             double* __restrict__ frame_stat) {
+  const int f = blockIdx.x * blockDim.x + threadIdx.x;
+  if (f >= frame_count) {
+    return;
+  }
+  double count = 0.0;
+  double sum = 0.0;
+  for (int slice = 0; slice < slice_count; ++slice) {
+    const double* entry =
+        slice_sum + (static_cast(f) * slice_count + slice) * 2;
+    count += entry[0];
+    sum += entry[1];
+  }
+  // A frame without a charged node has nothing to shift.
+  count = fmax(count, 1.0);
+  frame_stat[f * kStatWidth + kStatShift] = (sum - total_charge[f]) / count;
+  frame_stat[f * kStatWidth + kStatCount] = count;
+}
+
+}  // namespace
diff --git a/source/op/pt/long_range/latent_charge_sum.cu b/source/op/pt/long_range/latent_charge_sum.cu
new file mode 100644
index 0000000000..dea9bb1681
--- /dev/null
+++ b/source/op/pt/long_range/latent_charge_sum.cu
@@ -0,0 +1,254 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// CUDA entry of deepmd::latent_charge_sum (schema and formulas in ops.cc).
+//
+// A call evaluates its frames by the direct sum when they carry no cell and
+// by the mesh sum when they do (charge_sum.h). Across the domains of a
+// decomposed system, the ranks gather the positions and charges of the charged
+// nodes of every domain into one frame, in rank order (domain_exchange.h), and
+// each rank evaluates its own charged nodes against that frame; the outputs
+// return to the node axis of the rank, where the nodes it does not own carry
+// no charge.
+
+#include 
+#include 
+#include 
+
+#include 
+#include 
+#include 
+
+#include "charge_sum.h"
+#include "domain_exchange.h"
+
+namespace {
+
+using deepmd::long_range::ChargeSumOutputs;
+using deepmd::long_range::kLes;
+using deepmd::long_range::kMaxChannels;
+using deepmd::long_range::kMaxGaussians;
+using deepmd::long_range::kSog;
+
+// The evaluation of the frames, by the direct or the mesh sum, of the owner
+// nodes [owner_begin, owner_begin + owner_count), which are every node or lie
+// in a single frame. ``box`` is defined for periodic frames only. Without
+// nodes or frames the outputs are zero.
+struct Evaluation {
+  const torch::Tensor& box;
+  const torch::Tensor& coefficient;
+  int64_t kernel;
+  int64_t channel;
+  double prefactor;
+  int64_t order;
+  double spacing;
+  bool periodic;
+
+  ChargeSumOutputs operator()(const torch::Tensor& coord,
+                              const torch::Tensor& charge,
+                              const torch::Tensor& node_mask,
+                              const torch::Tensor& n_node_per_frame,
+                              const torch::Tensor& total_charge,
+                              int owner_begin,
+                              int owner_count,
+                              const torch::Tensor* communicator) const {
+    const int64_t frame_count = n_node_per_frame.size(0);
+    if (coord.size(0) == 0 || frame_count == 0) {
+      return {torch::zeros({owner_count, 1}, coord.options()),
+              torch::zeros({owner_count, charge.size(1)}, charge.options()),
+              torch::zeros({owner_count, 3}, charge.options()),
+              torch::zeros({owner_count, 9}, charge.options()),
+              torch::zeros({frame_count, 9}, charge.options())};
+    }
+    TORCH_CHECK(
+        (owner_begin == 0 && owner_count == coord.size(0)) || frame_count == 1,
+        "latent_charge_sum: owners other than every node lie in one "
+        "frame");
+    if (periodic) {
+      return deepmd::long_range::mesh_sum(
+          coord, charge, node_mask, n_node_per_frame, total_charge, box,
+          coefficient, kernel, channel, prefactor, order, spacing, owner_begin,
+          owner_count, communicator);
+    }
+    return deepmd::long_range::direct_sum(
+        coord, charge, node_mask, n_node_per_frame, total_charge, coefficient,
+        kernel, channel, prefactor, owner_begin, owner_count, communicator);
+  }
+};
+
+// Evaluates the frame of a system decomposed into the domains of
+// ``communicator``, of which this rank holds one: its charged nodes are its
+// owners, the charged nodes of every domain its partners. The nodes the rank
+// does not own carry no charge, and the per-frame virial is the share of its
+// owners.
+ChargeSumOutputs evaluate_across_domains(const Evaluation& evaluation,
+                                         const torch::Tensor& coord,
+                                         const torch::Tensor& charge,
+                                         const torch::Tensor& node_mask,
+                                         const torch::Tensor& n_node_per_frame,
+                                         const torch::Tensor& total_charge,
+                                         const torch::Tensor& communicator) {
+  TORCH_CHECK(n_node_per_frame.size(0) == 1,
+              "latent_charge_sum: an evaluation across MPI ranks takes one "
+              "frame, got ",
+              n_node_per_frame.size(0));
+  const int64_t node_count = coord.size(0);
+  // The charged nodes of the frame; the padding nodes past it carry none.
+  // Their count is the one value the host has to wait for before the
+  // exchange.
+  const auto in_frame =
+      torch::arange(node_count, n_node_per_frame.options()) < n_node_per_frame;
+  const auto owned = torch::nonzero(node_mask & in_frame).squeeze(1);
+  const auto rows =
+      torch::cat({coord.index_select(0, owned),
+                  charge.index_select(0, owned).to(torch::kFloat64)},
+                 1)
+          .contiguous();
+  const auto [partners, owner_begin] =
+      deepmd::long_range::gather_domain_rows(communicator, rows);
+  const int64_t partner_count = partners.size(0);
+  const auto partner_coord = partners.slice(1, 0, 3).contiguous();
+  const auto partner_charge =
+      partners.slice(1, 3).to(torch::kFloat32).contiguous();
+  const auto partner_mask = torch::ones({partner_count}, node_mask.options());
+  const auto partner_frame =
+      torch::full({1}, partner_count, n_node_per_frame.options());
+  const auto [owner_energy, owner_gradient, owner_force, owner_atom_virial,
+              virial] =
+      evaluation(partner_coord, partner_charge, partner_mask, partner_frame,
+                 total_charge, static_cast(owner_begin),
+                 static_cast(owned.size(0)), &communicator);
+  return {torch::zeros({node_count, 1}, coord.options())
+              .index_copy_(0, owned, owner_energy),
+          torch::zeros({node_count, charge.size(1)}, charge.options())
+              .index_copy_(0, owned, owner_gradient),
+          torch::zeros({node_count, 3}, charge.options())
+              .index_copy_(0, owned, owner_force),
+          torch::zeros({node_count, 9}, charge.options())
+              .index_copy_(0, owned, owner_atom_virial),
+          virial};
+}
+
+// Whether the frames carry cells: true when ``box`` is given and none of its
+// rows is zero, false when it is absent or all of its rows are; a call that
+// mixes the two is rejected.
+bool periodic_frames(const c10::optional& box) {
+  if (!box.has_value()) {
+    return false;
+  }
+  const auto cells = box->contiguous();
+  const double* value = cells.data_ptr();
+  const int64_t frame_count = cells.size(0);
+  int64_t periodic = 0;
+  for (int64_t f = 0; f < frame_count; ++f) {
+    periodic += std::any_of(value + 9 * f, value + 9 * (f + 1),
+                            [](double entry) { return entry != 0.0; });
+  }
+  TORCH_CHECK(periodic == 0 || periodic == frame_count,
+              "latent_charge_sum: every frame of a call is periodic, or none "
+              "is");
+  return periodic > 0;
+}
+
+}  // namespace
+
+namespace deepmd {
+namespace long_range {
+
+ChargeSumOutputs latent_charge_sum(torch::Tensor coord,
+                                   torch::Tensor charge,
+                                   torch::Tensor node_mask,
+                                   torch::Tensor n_node_per_frame,
+                                   torch::Tensor total_charge,
+                                   c10::optional box,
+                                   torch::Tensor coefficient,
+                                   torch::Tensor communicator,
+                                   int64_t kernel,
+                                   int64_t channel,
+                                   double prefactor,
+                                   int64_t order,
+                                   double spacing) {
+  TORCH_CHECK(coord.is_cuda() && charge.is_cuda() && node_mask.is_cuda() &&
+                  n_node_per_frame.is_cuda() && total_charge.is_cuda() &&
+                  coefficient.is_cuda(),
+              "latent_charge_sum: inputs must be CUDA tensors, except box and "
+              "communicator");
+  TORCH_CHECK(charge.device() == coord.device() &&
+                  node_mask.device() == coord.device() &&
+                  n_node_per_frame.device() == coord.device() &&
+                  total_charge.device() == coord.device() &&
+                  coefficient.device() == coord.device(),
+              "latent_charge_sum: inputs must share one device");
+  TORCH_CHECK(coord.dim() == 2 && coord.size(1) == 3 &&
+                  coord.scalar_type() == torch::kFloat64,
+              "latent_charge_sum: coord must be float64 with shape (N, 3)");
+  TORCH_CHECK(charge.dim() == 2 && charge.size(0) == coord.size(0) &&
+                  charge.scalar_type() == torch::kFloat32,
+              "latent_charge_sum: charge must be float32 with shape (N, C)");
+  TORCH_CHECK(node_mask.dim() == 1 && node_mask.size(0) == coord.size(0) &&
+                  node_mask.scalar_type() == torch::kBool,
+              "latent_charge_sum: node_mask must be bool with shape (N,)");
+  TORCH_CHECK(n_node_per_frame.dim() == 1,
+              "latent_charge_sum: n_node_per_frame must be one-dimensional");
+  TORCH_CHECK(total_charge.dim() == 1 &&
+                  total_charge.size(0) == n_node_per_frame.size(0) &&
+                  total_charge.scalar_type() == torch::kFloat64,
+              "latent_charge_sum: total_charge must be float64 with one entry "
+              "per frame");
+  TORCH_CHECK(
+      !box.has_value() ||
+          (box->device().is_cpu() && box->scalar_type() == torch::kFloat64 &&
+           box->dim() == 2 && box->size(0) == n_node_per_frame.size(0) &&
+           box->size(1) == 9),
+      "latent_charge_sum: box must be a float64 CPU tensor with shape "
+      "(F, 9)");
+  TORCH_CHECK(
+      coefficient.dim() == 1 && coefficient.scalar_type() == torch::kFloat64,
+      "latent_charge_sum: coefficient must be a float64 vector");
+  const int64_t channels = charge.size(1);
+  TORCH_CHECK(channels >= 1 && channels <= kMaxChannels,
+              "latent_charge_sum: the channel count must lie in [1, ",
+              kMaxChannels, "], got ", channels);
+  TORCH_CHECK(channel >= -1 && channel < channels,
+              "latent_charge_sum: the constrained channel must be -1 or lie "
+              "in [0, ",
+              channels, "), got ", channel);
+  if (kernel == kLes) {
+    TORCH_CHECK(coefficient.numel() == 1,
+                "latent_charge_sum: the erf kernel takes one coefficient");
+  } else if (kernel == kSog) {
+    TORCH_CHECK(coefficient.numel() % 2 == 0 && coefficient.numel() > 0 &&
+                    coefficient.numel() <= 2 * kMaxGaussians,
+                "latent_charge_sum: the Gaussian kernel takes two "
+                "coefficients per Gaussian, for at most ",
+                kMaxGaussians, " Gaussians");
+  } else {
+    TORCH_CHECK(false, "latent_charge_sum: unknown kernel ", kernel);
+  }
+  TORCH_CHECK(coord.size(0) < std::numeric_limits::max(),
+              "latent_charge_sum: too many nodes");
+
+  const c10::cuda::CUDAGuard guard(coord.device());
+  coord = coord.contiguous();
+  charge = charge.contiguous();
+  node_mask = node_mask.contiguous();
+  total_charge = total_charge.contiguous();
+  const bool periodic = periodic_frames(box);
+  const torch::Tensor cells = periodic ? box->contiguous() : torch::Tensor();
+  coefficient = coefficient.contiguous();
+  const Evaluation evaluation{cells,     coefficient, kernel,  channel,
+                              prefactor, order,       spacing, periodic};
+  if (deepmd::long_range::domain_count(communicator) == 1) {
+    return evaluation(coord, charge, node_mask, n_node_per_frame, total_charge,
+                      0, static_cast(coord.size(0)), nullptr);
+  }
+  return evaluate_across_domains(evaluation, coord, charge, node_mask,
+                                 n_node_per_frame, total_charge, communicator);
+}
+
+}  // namespace long_range
+}  // namespace deepmd
+
+TORCH_LIBRARY_FRAGMENT(deepmd, library) {
+  library.impl("latent_charge_sum", torch::kCUDA,
+               &deepmd::long_range::latent_charge_sum);
+}
diff --git a/source/op/pt/long_range/mesh_sum.cu b/source/op/pt/long_range/mesh_sum.cu
new file mode 100644
index 0000000000..3fa09ef8c2
--- /dev/null
+++ b/source/op/pt/long_range/mesh_sum.cu
@@ -0,0 +1,1544 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// CUDA kernels of the latent-charge long range of periodic frames, by the
+// smooth particle-mesh Ewald method (schema and formulas in ops.cc, the
+// evaluation in deepmd/dpmodel/long_range/reciprocal.py).
+//
+// Every frame owns a mesh of n1 x n2 x n3 points spanning its cell. The
+// charges spread to the mesh with cardinal B-splines of order p along the
+// fractional coordinates s = r B^-1, the lattice vectors being the rows of the
+// cell B: the scaled coordinate u_d = n_d s_d reaches the p points
+// floor(u_d) - p + 1, ..., floor(u_d) with the weights M_p(u_d - m). A real
+// FFT gives the half spectrum Q(m) of the charge mesh, and the influence
+// function
+//
+//   G(m) = prefactor * K^(k_m) / V * prod_d |D_d(m_d)|^-2,  k_m = 2 pi B^-1 m,
+//
+// with the Fourier transform K^ of the kernel and the assignment factor
+// D_d(m) = sum_l M_p(l + 1) exp(2 pi i m l / n_d), turns it into the
+// potential mesh, the unnormalized inverse transform of G Q. The potential at
+// a node is its interpolation by the B-splines of the node, less the potential
+// prefactor * K(0) q of the charge at itself, and the force at fixed charges
+// the negative gradient of that interpolation times the charges. The zero mode
+// and the modes at the Nyquist frequency of an axis, which have no unique wave
+// number, carry no weight.
+//
+// The virial of a frame is that of the strain of its cell, under which the
+// fractional coordinates stay put while k_m and the volume change:
+//
+//   W = sum_m e_m (1 + 2 dln K^ / d(k^2) k_m (x) k_m),  e_m = G |Q|^2 / 2,
+//
+// the sum over the full spectrum, every interior plane of the half spectrum
+// counting twice. Its per-node share is F_i (x) r_i, the force index first,
+// plus the rest of W shared evenly over the charged nodes.
+//
+// The charges spread in 64-bit fixed point, so that the mesh, and every
+// output, does not depend on the order of the atomic additions. The sums over
+// the modes and over the nodes of a frame add float64 sums in a fixed order.
+// The meshes, the transforms and the influence function are float32, against
+// a float64 throughput of a small fraction of the float32 one on many
+// devices; the B-spline weights are evaluated in float64, whose rounding to
+// float32 leaves their sum unbiased. The forces then carry the rounding of
+// float32, and the energy of a frame that of the reciprocal sum including the
+// interaction of every charge with itself, which the self term cancels: a few
+// 1e-7 of the self energy.
+//
+// The nodes whose outputs a call evaluates, its owners, are either every node
+// or, across the domains of a decomposed system, the charged nodes of one MPI
+// rank, whose frame then holds the charged nodes of every rank. Every rank
+// spreads every charge and transforms the whole mesh, and interpolates only at
+// its owners; the per-frame sums of the owners are added over the domains.
+
+#include 
+#include 
+#include 
+#include 
+#include 
+
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+
+#include "charge_sum.h"
+#include "frame_sums.cuh"
+
+namespace {
+
+constexpr int kThreads = 256;
+constexpr int kMaxOrder = 12;
+
+using deepmd::long_range::kLes;
+using deepmd::long_range::kMaxGaussians;
+using deepmd::long_range::kSog;
+
+// Resolution of the fixed-point charge mesh: 2^-40 e, far below the float32
+// precision of the raw charges, for a sum of at most 2^23 e at a point.
+constexpr float kFixedScaleFloat = 1099511627776.f;  // 2^40
+
+constexpr double kPi = 3.14159265358979323846;
+
+// The cell of one frame and its mesh. The host sizes the mesh, which fixes
+// the buffers and the launches of a call, and the device derives the
+// geometry from the cell and the prefactor of the call, so that a CUDA graph
+// capturing a call evaluates the cell of every replay on the meshes of the
+// capture.
+struct FrameMesh {
+  // B^-1, row-major: s_d = sum_j r_j inverse[3 j + d].
+  double inverse[9];
+  // 2 pi B^-1 in float32, row-major like inverse: the wave vector of the mode
+  // m is k_a = sum_d reciprocal[3 a + d] m_d.
+  float reciprocal[9];
+  // The factor prefactor / V of the influence function in float32, V the
+  // volume of the cell.
+  float influence_scale;
+  // -n_d (B^-1)_jd, row-major in (j, d): the force at fixed charges from the
+  // derivatives of the energy with respect to the scaled coordinates u_d.
+  float gradient_to_force[9];
+  int n[3];
+  // The first point of the meshes of the frame, channel-major, in the mesh
+  // buffer of the call.
+  int64_t offset;
+};
+
+// The smallest even number of at least ``n`` without a prime factor above 7.
+int64_t even_fft_size(int64_t n) {
+  int64_t m = std::max(2, n + n % 2);
+  while (true) {
+    int64_t rest = m;
+    for (const int64_t p : {2, 3, 5, 7}) {
+      while (rest % p == 0) {
+        rest /= p;
+      }
+    }
+    if (rest == 1) {
+      return m;
+    }
+    m += 2;
+  }
+}
+
+// The mesh of the cell ``cell`` at a spacing of at most ``spacing``, sized
+// along every lattice vector; store_frame_mesh adds the geometry of the cell
+// on the device.
+FrameMesh frame_mesh_size(const double* cell, double spacing) {
+  FrameMesh mesh{};
+  const double a = cell[0], b = cell[1], c = cell[2];
+  const double d = cell[3], e = cell[4], f = cell[5];
+  const double g = cell[6], h = cell[7], i = cell[8];
+  const double det =
+      a * (e * i - f * h) - b * (d * i - f * g) + c * (d * h - e * g);
+  // A volume that is a vanishing fraction of the product of the lengths of
+  // the lattice vectors, or not finite, spans no lattice.
+  const double lengths = std::sqrt(a * a + b * b + c * c) *
+                         std::sqrt(d * d + e * e + f * f) *
+                         std::sqrt(g * g + h * h + i * i);
+  TORCH_CHECK(std::abs(det) > 1e-6 * lengths,
+              "latent_charge_sum: a periodic cell is singular");
+  // The squared length of a lattice vector sums the squares of its
+  // components, each rounded on its own as the Python front end and the C++
+  // interface form them: the explicit fma of a square with zero is that
+  // rounded square, which the compiler does not fuse into the sum, so that
+  // all three size the mesh identically.
+  for (int axis = 0; axis < 3; ++axis) {
+    const double* v = cell + 3 * axis;
+    const double length =
+        std::sqrt(std::fma(v[0], v[0], 0.0) + std::fma(v[1], v[1], 0.0) +
+                  std::fma(v[2], v[2], 0.0));
+    const double points = std::ceil(length / spacing);
+    TORCH_CHECK(points < 1 << 20,
+                "latent_charge_sum: the mesh of a cell exceeds 2^20 points "
+                "along a lattice vector");
+    mesh.n[axis] =
+        static_cast(even_fft_size(static_cast(points)));
+  }
+  TORCH_CHECK(int64_t{mesh.n[0]} * mesh.n[1] * mesh.n[2] <
+                  std::numeric_limits::max(),
+              "latent_charge_sum: the mesh of a cell exceeds 2^31 points");
+  return mesh;
+}
+
+__host__ __device__ int64_t mesh_points(const FrameMesh& mesh) {
+  return int64_t{mesh.n[0]} * mesh.n[1] * mesh.n[2];
+}
+
+// Nodes per block of the kernels that spread and interpolate, one warp each.
+constexpr int kStencilWarps = 4;
+constexpr int kWarp = 32;
+
+// The stencil of one node along one axis of its mesh: for each of its p
+// points on the axis, its offset on the flat mesh of a frame along that axis,
+// the index of the point times the stride of the axis, and its B-spline
+// weight and slope, evaluated in float64 and stored in float32, the precision
+// of the outputs: rounding a recursion of float32 steps instead would bias
+// the sum of the weights.
+template 
+struct AxisStencil {
+  int offset[kOrder];
+  float value[kOrder];
+  float slope[kOrder];
+};
+
+// The stencil of a node along axis ``d``. The scaled coordinate u = n_d s_d
+// of the fractional coordinate s_d reaches the points floor(u) - p + 1 + k,
+// k = 0, ..., p - 1, with the weights value[k] = M_p(t + p - 1 - k), t the
+// fractional part of u, and slope[k] their derivative with respect to u. The
+// recursion M_n(t + j) = ((t + j) M_{n-1}(t + j) + (n - t - j) M_{n-1}(t + j -
+// 1)) / (n - 1) builds a_j = M_n(t + j) from M_2; the slope follows from
+// M_p'(x) = M_{p-1}(x) - M_{p-1}(x - 1).
+template 
+__device__ void axis_stencil(const double* coord,
+                             const FrameMesh& mesh,
+                             int d,
+                             AxisStencil& stencil) {
+  double s = coord[0] * mesh.inverse[d] + coord[1] * mesh.inverse[3 + d] +
+             coord[2] * mesh.inverse[6 + d];
+  s -= floor(s);
+  const double u = s * mesh.n[d];
+  const double base = floor(u);
+  const double t = u - base;
+  double a[kOrder];
+  a[0] = t;
+  a[1] = 1.0 - t;
+#pragma unroll
+  for (int n = 3; n <= kOrder; ++n) {
+    if (n == kOrder) {
+      // The weights of order p - 1 give the slopes.
+#pragma unroll
+      for (int j = 0; j < kOrder; ++j) {
+        const double high = j < kOrder - 1 ? a[j] : 0.0;
+        const double low = j > 0 ? a[j - 1] : 0.0;
+        stencil.slope[kOrder - 1 - j] = static_cast(high - low);
+      }
+    }
+    const double scale = 1.0 / (n - 1);
+    a[n - 1] = scale * (1.0 - t) * a[n - 2];
+#pragma unroll
+    for (int j = n - 2; j > 0; --j) {
+      a[j] = scale * ((t + j) * a[j] + (n - t - j) * a[j - 1]);
+    }
+    a[0] = scale * t * a[0];
+  }
+  if constexpr (kOrder == 2) {
+    stencil.slope[0] = -1.f;
+    stencil.slope[1] = 1.f;
+  }
+  const int n = mesh.n[d];
+  const int stride = d == 0 ? mesh.n[1] * mesh.n[2] : d == 1 ? mesh.n[2] : 1;
+  // A stencil wider than the axis wraps more than once.
+  const int first = ((static_cast(base) - kOrder + 1) % n + n) % n;
+#pragma unroll
+  for (int k = 0; k < kOrder; ++k) {
+    stencil.value[k] = static_cast(a[kOrder - 1 - k]);
+    stencil.offset[k] = (first + k) % n * stride;
+  }
+}
+
+// Writes, one thread per node and axis, the stencils of the charged nodes
+// along the three axes of their meshes, and, on the thread of the first axis,
+// the charges of every node, shifted in the constrained channel and zero on
+// the nodes that carry none. Every float64 warp instruction occupies the
+// float64 pipe of the device whatever the number of its active lanes, so the
+// float64 work of a node runs here with every lane busy rather than in the
+// warps of the kernels that spread and interpolate.
+template 
+__global__ void __launch_bounds__(kThreads)
+    node_stencils(const double* __restrict__ coord,
+                  const float* __restrict__ charge,
+                  const bool* __restrict__ node_mask,
+                  const int64_t* __restrict__ frame_ptr,
+                  const FrameMesh* __restrict__ frames,
+                  const double* __restrict__ frame_stat,
+                  int frame_count,
+                  int64_t node_count,
+                  int channels,
+                  int channel,
+                  AxisStencil* __restrict__ stencils,
+                  float* __restrict__ node_charge) {
+  for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x;
+       index < 3 * node_count; index += int64_t{gridDim.x} * blockDim.x) {
+    const int64_t node = index / 3;
+    const int axis = static_cast(index - 3 * node);
+    const int f = frame_of(frame_ptr, frame_count, node);
+    const bool charged = f >= 0 && node_mask[node];
+    if (charged) {
+      axis_stencil(coord + 3 * node, frames[f], axis, stencils[index]);
+    }
+    if (axis == 0) {
+      const double shift = charged && channel >= 0
+                               ? frame_stat[f * kStatWidth + kStatShift]
+                               : 0.0;
+      for (int c = 0; c < channels; ++c) {
+        const double raw = charged ? charge[node * channels + c] : 0.0;
+        node_charge[node * channels + c] =
+            static_cast(c == channel ? raw - shift : raw);
+      }
+    }
+  }
+}
+
+// Rounds ``x`` to the nearest integer, ties to even, as __float2ll_rn does,
+// in integer operations: the conversion instruction occupies the float64 pipe.
+// A finite x is a 24-bit significand times 2^exponent; a negative exponent
+// shifts the significand right with rounding. Magnitudes of 2^63 and beyond,
+// infinities and NaN saturate.
+__device__ __forceinline__ long long float_to_fixed(float x) {
+  const unsigned bits = __float_as_uint(x);
+  const int biased = static_cast((bits >> 23) & 0xffu);
+  // |x| < 1/2, zero and subnormal numbers included, rounds to zero.
+  if (biased < 126) {
+    return 0;
+  }
+  const int exponent = biased - 150;
+  if (exponent >= 40) {
+    return bits >> 31 ? std::numeric_limits::min()
+                      : std::numeric_limits::max();
+  }
+  const long long significand = (bits & 0x7fffffu) | 0x800000u;
+  long long magnitude;
+  if (exponent >= 0) {
+    magnitude = significand << exponent;
+  } else {
+    const int shift = -exponent;
+    const long long kept = significand >> shift;
+    const long long rest = significand & ((1ll << shift) - 1);
+    const long long half = 1ll << (shift - 1);
+    magnitude = kept + (rest > half || (rest == half && (kept & 1)) ? 1 : 0);
+  }
+  return bits >> 31 ? -magnitude : magnitude;
+}
+
+// The assignment of the p^3 points of a stencil to the lanes of a warp: the
+// lanes take the (k2, k3) pairs, k3 running fastest, so that a warp touches
+// contiguous runs of the mesh, and every lane walks the k1 axis for its
+// pairs; while the pairs are fewer than the lanes, groups of lanes split the
+// k1 axis.
+template 
+struct LaneSlots {
+  static constexpr int kPairs = kOrder * kOrder;
+  static constexpr int kGroups = kPairs >= kWarp ? 1 : kWarp / kPairs;
+  static constexpr int kSlots = kGroups * kPairs;
+  static constexpr int kRounds = (kSlots + kWarp - 1) / kWarp;
+};
+
+// Copies the stencils of a node along its three axes to the shared memory of
+// its warp.
+template 
+__device__ __forceinline__ void load_stencil(
+    const AxisStencil* __restrict__ stencils,
+    int64_t node,
+    int lane,
+    AxisStencil* shared) {
+  constexpr int kWords = 3 * sizeof(AxisStencil) / sizeof(int);
+  const int* source = reinterpret_cast(stencils + 3 * node);
+  int* target = reinterpret_cast(shared);
+  for (int word = lane; word < kWords; word += kWarp) {
+    target[word] = source[word];
+  }
+  __syncwarp();
+}
+
+// Spreads the charges of the charged nodes to the meshes of their frames, one
+// warp per node, one pass per channel.
+template 
+__global__ void __launch_bounds__(kStencilWarps* kWarp)
+    spread_charges(const bool* __restrict__ node_mask,
+                   const int64_t* __restrict__ frame_ptr,
+                   const FrameMesh* __restrict__ frames,
+                   const AxisStencil* __restrict__ stencils,
+                   const float* __restrict__ node_charge,
+                   int frame_count,
+                   int64_t node_count,
+                   int channels,
+                   unsigned long long* __restrict__ mesh) {
+  using Slots = LaneSlots;
+  __shared__ AxisStencil shared[kStencilWarps][3];
+  const int warp = threadIdx.x / kWarp;
+  const int lane = threadIdx.x % kWarp;
+  const AxisStencil* axes = shared[warp];
+  for (int64_t node = int64_t{blockIdx.x} * kStencilWarps + warp;
+       node < node_count; node += int64_t{gridDim.x} * kStencilWarps) {
+    const int f = frame_of(frame_ptr, frame_count, node);
+    if (f < 0 || !node_mask[node]) {
+      continue;
+    }
+    load_stencil(stencils, node, lane, shared[warp]);
+    const FrameMesh& grid = frames[f];
+    for (int c = 0; c < channels; ++c) {
+      const float q = node_charge[node * channels + c] * kFixedScaleFloat;
+      unsigned long long* channel_mesh =
+          mesh + grid.offset + c * mesh_points(grid);
+#pragma unroll
+      for (int round = 0; round < Slots::kRounds; ++round) {
+        const int slot = lane + kWarp * round;
+        if (slot < Slots::kSlots) {
+          const int group = slot / Slots::kPairs;
+          const int pair = slot - group * Slots::kPairs;
+          const int k2 = pair / kOrder;
+          const int k3 = pair - k2 * kOrder;
+          const float weight = q * axes[1].value[k2] * axes[2].value[k3];
+          const int column = axes[1].offset[k2] + axes[2].offset[k3];
+#pragma unroll
+          for (int k1 = group; k1 < kOrder; k1 += Slots::kGroups) {
+            atomicAdd(channel_mesh + axes[0].offset[k1] + column,
+                      static_cast(
+                          float_to_fixed(weight * axes[0].value[k1])));
+          }
+        }
+      }
+    }
+    // The stencil of the next node of the warp overwrites this one.
+    __syncwarp();
+  }
+}
+
+// Stores the mesh of one frame, sized on the host and carried by value, with
+// the geometry of its cell ``cell``. Thread k = 3 j + d of kGeometryThreads
+// forms the entry (B^-1)_jd = C_dj / det B of the inverse of the cell B, C
+// its cofactors, and the wave-vector and force factors of that entry; the
+// first thread also forms the factor prefactor / V of the influence
+// function.
+constexpr int kGeometryThreads = 9;
+
+__global__ void store_frame_mesh(FrameMesh mesh,
+                                 const double* __restrict__ cell,
+                                 double prefactor,
+                                 FrameMesh* __restrict__ out) {
+  const int k = threadIdx.x;
+  const int j = k / 3;
+  const int d = k % 3;
+  const double det = cell[0] * (cell[4] * cell[8] - cell[5] * cell[7]) -
+                     cell[1] * (cell[3] * cell[8] - cell[5] * cell[6]) +
+                     cell[2] * (cell[3] * cell[7] - cell[4] * cell[6]);
+  const double* row1 = cell + 3 * ((d + 1) % 3);
+  const double* row2 = cell + 3 * ((d + 2) % 3);
+  const double inverse = (row1[(j + 1) % 3] * row2[(j + 2) % 3] -
+                          row1[(j + 2) % 3] * row2[(j + 1) % 3]) /
+                         det;
+  out->inverse[k] = inverse;
+  out->reciprocal[k] = static_cast(2.0 * kPi * inverse);
+  out->gradient_to_force[k] = static_cast(-mesh.n[d] * inverse);
+  if (k == 0) {
+    out->influence_scale = static_cast(prefactor / fabs(det));
+    for (int axis = 0; axis < 3; ++axis) {
+      out->n[axis] = mesh.n[axis];
+    }
+    out->offset = mesh.offset;
+  }
+}
+
+// Converts the fixed-point mesh to the float32 mesh that the transforms read.
+__global__ void fixed_to_float(int64_t count,
+                               const unsigned long long* __restrict__ fixed,
+                               float* __restrict__ mesh) {
+  constexpr float kFixedInverse = 1.f / kFixedScaleFloat;
+  for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x;
+       index < count; index += int64_t{gridDim.x} * blockDim.x) {
+    mesh[index] =
+        __ll2float_rn(static_cast(fixed[index])) * kFixedInverse;
+  }
+}
+
+// The weights M_p(1), ..., M_p(p - 1) of the assignment factor.
+struct IntegerWeights {
+  double value[kMaxOrder - 1];
+};
+
+// The inverse squared assignment factor |D_d(m)|^-2 along every axis of the
+// half spectrum of a mesh: n1 entries, then n2, then n3 / 2 + 1. It depends
+// only on the shape of the mesh and the order of the B-splines, so the plans
+// of a shape hold it (create_mesh_plans).
+__global__ void assignment_factor(FrameMesh mesh,
+                                  int order,
+                                  IntegerWeights integer_weight,
+                                  float* __restrict__ factor) {
+  const int half = mesh.n[2] / 2 + 1;
+  const int total = mesh.n[0] + mesh.n[1] + half;
+  for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < total;
+       index += gridDim.x * blockDim.x) {
+    int axis = 0;
+    int m = index;
+    if (m >= mesh.n[0]) {
+      m -= mesh.n[0];
+      axis = 1;
+      if (m >= mesh.n[1]) {
+        m -= mesh.n[1];
+        axis = 2;
+      }
+    }
+    const double theta = 2.0 * kPi * m / mesh.n[axis];
+    double re = 0.0;
+    double im = 0.0;
+    for (int l = 0; l < order - 1; ++l) {
+      double s;
+      double c;
+      sincos(theta * l, &s, &c);
+      re += integer_weight.value[l] * c;
+      im += integer_weight.value[l] * s;
+    }
+    factor[index] = static_cast(1.0 / (re * re + im * im));
+  }
+}
+
+// The terms of the transform of the kernel, in float32, the precision of the
+// transforms: 1 / (4 alpha^2) of the les kernel, or, for every Gaussian of the
+// sog kernel, its weight a_g (2 pi)^(3/2) sigma_g^3 and half its squared
+// width, from ``coefficient``, which holds alpha, or the amplitudes and then
+// the rates c_g = log2(e) / (2 sigma_g^2).
+struct KernelTerms {
+  float inv_four_alpha_sq;
+  float weight[kMaxGaussians];
+  float half_width_sq[kMaxGaussians];
+};
+
+__device__ void load_kernel_terms(int kernel,
+                                  const double* __restrict__ coefficient,
+                                  int gaussians,
+                                  KernelTerms& terms) {
+  if (kernel == kLes) {
+    if (threadIdx.x == 0) {
+      const double alpha = coefficient[0];
+      terms.inv_four_alpha_sq = static_cast(0.25 / (alpha * alpha));
+    }
+    return;
+  }
+  // (2 pi)^{3/2}
+  constexpr double kGaussianNorm = 15.749609945722419;
+  for (int g = threadIdx.x; g < gaussians; g += blockDim.x) {
+    // sigma^2 = log2(e) / (2 c) = 1 / (2 c ln 2)
+    const double width_sq = 0.72134752044448170 / coefficient[gaussians + g];
+    terms.weight[g] = static_cast(coefficient[g] * kGaussianNorm *
+                                         width_sq * sqrt(width_sq));
+    terms.half_width_sq[g] = static_cast(0.5 * width_sq);
+  }
+}
+
+// The transform of the kernel at k^2 and its derivative with respect to k^2.
+__device__ __forceinline__ void kernel_transform(int kernel,
+                                                 const KernelTerms& terms,
+                                                 int gaussians,
+                                                 float k_sq,
+                                                 float& value,
+                                                 float& slope) {
+  if (kernel == kLes) {
+    value = 4.f * static_cast(kPi) *
+            expf(-k_sq * terms.inv_four_alpha_sq) / k_sq;
+    slope = value * (-terms.inv_four_alpha_sq - 1.f / k_sq);
+    return;
+  }
+  value = 0.f;
+  slope = 0.f;
+  for (int g = 0; g < gaussians; ++g) {
+    const float term = terms.weight[g] * expf(-terms.half_width_sq[g] * k_sq);
+    value += term;
+    slope -= terms.half_width_sq[g] * term;
+  }
+}
+
+// The influence function of one frame on its half spectrum: multiplies every
+// mode of the spectra of the channels by G in place and sums, per block, the
+// energy of the modes and their virial, xx yy zz xy xz yz. The modes are
+// evaluated in float32, the precision of the transforms, summed in float32
+// over the modes of a row that a lane takes, and in float64 over the rows.
+constexpr int kModeSums = 7;
+
+__global__ void __launch_bounds__(kThreads)
+    apply_influence(const FrameMesh* __restrict__ frame,
+                    int channels,
+                    int kernel,
+                    const double* __restrict__ coefficient,
+                    int gaussians,
+                    const float* __restrict__ factor,
+                    float2* __restrict__ spectrum,
+                    double* __restrict__ block_sum) {
+  __shared__ KernelTerms terms;
+  load_kernel_terms(kernel, coefficient, gaussians, terms);
+  __syncthreads();
+  const FrameMesh& mesh = *frame;
+  const int n1 = mesh.n[0];
+  const int n2 = mesh.n[1];
+  const int n3 = mesh.n[2];
+  const int half = n3 / 2 + 1;
+  const int modes = n1 * n2 * half;
+  // k_a = sum_d reciprocal[3 a + d] m_d = 2 pi sum_d (B^-1)_ad m_d
+  float reciprocal[9];
+#pragma unroll
+  for (int k = 0; k < 9; ++k) {
+    reciprocal[k] = mesh.reciprocal[k];
+  }
+  const float scale = mesh.influence_scale;
+  // Every warp takes rows (i1, i2) of the half spectrum, its lanes the modes
+  // of a row, which lie contiguously.
+  const int lane = threadIdx.x % kWarp;
+  const int warps = gridDim.x * (kThreads / kWarp);
+  double sum[kModeSums] = {};
+  for (int row = blockIdx.x * (kThreads / kWarp) + threadIdx.x / kWarp;
+       row < n1 * n2; row += warps) {
+    float row_sum[kModeSums] = {};
+    const int i1 = row / n2;
+    const int i2 = row - i1 * n2;
+    const int m1 = i1 <= n1 / 2 ? i1 : i1 - n1;
+    const int m2 = i2 <= n2 / 2 ? i2 : i2 - n2;
+    const bool row_weightless = 2 * m1 == n1 || 2 * m2 == n2;
+    float k_row[3];
+#pragma unroll
+    for (int a = 0; a < 3; ++a) {
+      k_row[a] = reciprocal[3 * a] * m1 + reciprocal[3 * a + 1] * m2;
+    }
+    const float row_factor = scale * factor[i1] * factor[n1 + i2];
+    for (int i3 = lane; i3 < half; i3 += kWarp) {
+      const int m3 = i3;
+      const bool weightless =
+          row_weightless || 2 * m3 == n3 || (m1 == 0 && m2 == 0 && m3 == 0);
+      float influence = 0.f;
+      float influence_slope = 0.f;
+      float k[3] = {0.f, 0.f, 0.f};
+      if (!weightless) {
+#pragma unroll
+        for (int a = 0; a < 3; ++a) {
+          k[a] = k_row[a] + reciprocal[3 * a + 2] * m3;
+        }
+        const float k_sq = k[0] * k[0] + k[1] * k[1] + k[2] * k[2];
+        float transform;
+        float slope;
+        kernel_transform(kernel, terms, gaussians, k_sq, transform, slope);
+        influence = row_factor * transform * factor[n1 + n2 + i3];
+        influence_slope = row_factor * slope * factor[n1 + n2 + i3];
+      }
+      // An interior plane of the half spectrum stands for its mirror as well.
+      const float multiplicity = m3 == 0 ? 1.f : 2.f;
+      const int index = row * half + i3;
+      float power = 0.f;
+      for (int c = 0; c < channels; ++c) {
+        float2& value = spectrum[int64_t{c} * modes + index];
+        power += value.x * value.x + value.y * value.y;
+        value.x *= influence;
+        value.y *= influence;
+      }
+      const float energy = power * (0.5f * multiplicity * influence);
+      // Twice the derivative of the energy of the mode with respect to k^2,
+      // formed from the derivative of the transform, which stays defined
+      // where the transform vanishes.
+      const float twice = power * (multiplicity * influence_slope);
+      row_sum[0] += energy;
+      row_sum[1] += energy + twice * k[0] * k[0];
+      row_sum[2] += energy + twice * k[1] * k[1];
+      row_sum[3] += energy + twice * k[2] * k[2];
+      row_sum[4] += twice * k[0] * k[1];
+      row_sum[5] += twice * k[0] * k[2];
+      row_sum[6] += twice * k[1] * k[2];
+    }
+#pragma unroll
+    for (int s = 0; s < kModeSums; ++s) {
+      sum[s] += row_sum[s];
+    }
+  }
+  __shared__ double partial[kModeSums][kThreads];
+  for (int s = 0; s < kModeSums; ++s) {
+    partial[s][threadIdx.x] = sum[s];
+  }
+  __syncthreads();
+  for (int stride = kThreads / 2; stride > 0; stride /= 2) {
+    if (threadIdx.x < stride) {
+      for (int s = 0; s < kModeSums; ++s) {
+        partial[s][threadIdx.x] += partial[s][threadIdx.x + stride];
+      }
+    }
+    __syncthreads();
+  }
+  if (threadIdx.x < kModeSums) {
+    block_sum[blockIdx.x * kModeSums + threadIdx.x] = partial[threadIdx.x][0];
+  }
+}
+
+// Adds the block sums of apply_influence into the virial of the frame,
+// row-major: every thread sums a fixed strided subset of the blocks, and the
+// block adds the thread sums in a fixed order.
+__global__ void __launch_bounds__(kThreads)
+    mode_virial(const double* __restrict__ block_sum,
+                int block_count,
+                double* __restrict__ virial) {
+  double value[kModeSums] = {};
+  for (int block = threadIdx.x; block < block_count; block += kThreads) {
+#pragma unroll
+    for (int s = 0; s < kModeSums; ++s) {
+      value[s] += block_sum[block * kModeSums + s];
+    }
+  }
+  __shared__ double partial[kModeSums][kThreads];
+#pragma unroll
+  for (int s = 0; s < kModeSums; ++s) {
+    partial[s][threadIdx.x] = value[s];
+  }
+  __syncthreads();
+  for (int stride = kThreads / 2; stride > 0; stride /= 2) {
+    if (threadIdx.x < stride) {
+#pragma unroll
+      for (int s = 0; s < kModeSums; ++s) {
+        partial[s][threadIdx.x] += partial[s][threadIdx.x + stride];
+      }
+    }
+    __syncthreads();
+  }
+  // Components 1..6 are xx yy zz xy xz yz; the energy is not stored.
+  constexpr int kEntry[kModeSums][2] = {{-1, -1}, {0, 0}, {4, 4}, {8, 8},
+                                        {1, 3},   {2, 6}, {5, 7}};
+  if (threadIdx.x > 0 && threadIdx.x < kModeSums) {
+    virial[kEntry[threadIdx.x][0]] = partial[threadIdx.x][0];
+    virial[kEntry[threadIdx.x][1]] = partial[threadIdx.x][0];
+  }
+}
+
+// The contact value K(0) of the kernel.
+__device__ double contact_value(int kernel,
+                                const double* __restrict__ coefficient,
+                                int gaussians) {
+  if (kernel == kLes) {
+    return 1.1283791670955126 * coefficient[0];
+  }
+  double value = 0.0;
+  for (int g = 0; g < gaussians; ++g) {
+    value += coefficient[g];
+  }
+  return value;
+}
+
+// Adds ``value`` over the lanes of a warp in a fixed butterfly order; every
+// lane receives the total.
+__device__ __forceinline__ float warp_sum(float value) {
+#pragma unroll
+  for (int stride = kWarp / 2; stride > 0; stride /= 2) {
+    value += __shfl_xor_sync(0xffffffffu, value, stride);
+  }
+  return value;
+}
+
+// Interpolates the potential meshes at every owner node, one warp per node,
+// one pass per channel: the potential phi^rec_c of the mesh, which still holds
+// that of the charge at itself, and the force at fixed charges. Owners that
+// carry no charge get zeros. Every lane sums its points in float32, the
+// precision of the meshes, and the warp adds the lane sums.
+template 
+__global__ void __launch_bounds__(kStencilWarps* kWarp)
+    interpolate_owners(const bool* __restrict__ node_mask,
+                       const int64_t* __restrict__ frame_ptr,
+                       const FrameMesh* __restrict__ frames,
+                       const AxisStencil* __restrict__ stencils,
+                       const float* __restrict__ node_charge,
+                       const float* __restrict__ potential_mesh,
+                       int frame_count,
+                       int owner_begin,
+                       int owner_count,
+                       int channels,
+                       float* __restrict__ potential,
+                       float* __restrict__ force) {
+  using Slots = LaneSlots;
+  __shared__ AxisStencil shared[kStencilWarps][3];
+  const int warp = threadIdx.x / kWarp;
+  const int lane = threadIdx.x % kWarp;
+  const AxisStencil* axes = shared[warp];
+  for (int owner = blockIdx.x * kStencilWarps + warp; owner < owner_count;
+       owner += gridDim.x * kStencilWarps) {
+    const int64_t node = int64_t{owner_begin} + owner;
+    const int f = frame_of(frame_ptr, frame_count, node);
+    if (f < 0 || !node_mask[node]) {
+      for (int c = lane; c < channels; c += kWarp) {
+        potential[int64_t{owner} * channels + c] = 0.f;
+      }
+      if (lane < 3) {
+        force[int64_t{owner} * 3 + lane] = 0.f;
+      }
+      continue;
+    }
+    load_stencil(stencils, node, lane, shared[warp]);
+    const FrameMesh& grid = frames[f];
+    // The derivative of sum_c q_c phi_c with respect to u_d.
+    float gradient[3] = {0.f, 0.f, 0.f};
+    for (int c = 0; c < channels; ++c) {
+      const float* channel_mesh =
+          potential_mesh + grid.offset + c * mesh_points(grid);
+      float phi = 0.f;
+      float slope[3] = {0.f, 0.f, 0.f};
+#pragma unroll
+      for (int round = 0; round < Slots::kRounds; ++round) {
+        const int slot = lane + kWarp * round;
+        if (slot < Slots::kSlots) {
+          const int group = slot / Slots::kPairs;
+          const int pair = slot - group * Slots::kPairs;
+          const int k2 = pair / kOrder;
+          const int k3 = pair - k2 * kOrder;
+          const float v2 = axes[1].value[k2];
+          const float v3 = axes[2].value[k3];
+          const float v23 = v2 * v3;
+          const float s2v3 = axes[1].slope[k2] * v3;
+          const float v2s3 = v2 * axes[2].slope[k3];
+          const int column = axes[1].offset[k2] + axes[2].offset[k3];
+#pragma unroll
+          for (int k1 = group; k1 < kOrder; k1 += Slots::kGroups) {
+            const float value = channel_mesh[axes[0].offset[k1] + column];
+            const float v1_value = axes[0].value[k1] * value;
+            phi += v1_value * v23;
+            slope[0] += axes[0].slope[k1] * value * v23;
+            slope[1] += v1_value * s2v3;
+            slope[2] += v1_value * v2s3;
+          }
+        }
+      }
+      phi = warp_sum(phi);
+      const float q = node_charge[node * channels + c];
+#pragma unroll
+      for (int d = 0; d < 3; ++d) {
+        gradient[d] += q * warp_sum(slope[d]);
+      }
+      if (lane == 0) {
+        potential[int64_t{owner} * channels + c] = phi;
+      }
+    }
+    if (lane < 3) {
+      // F_j = -sum_d dE/du_d n_d (B^-1)_jd
+      force[int64_t{owner} * 3 + lane] =
+          gradient[0] * grid.gradient_to_force[3 * lane] +
+          gradient[1] * grid.gradient_to_force[3 * lane + 1] +
+          gradient[2] * grid.gradient_to_force[3 * lane + 2];
+    }
+    // The stencil of the next node of the warp overwrites this one.
+    __syncwarp();
+  }
+}
+
+// The potential prefactor * K(0) of a unit charge at itself, which the
+// potential mesh holds and the potential of a node excludes, for every thread
+// of the block.
+__device__ double self_potential(int kernel,
+                                 const double* __restrict__ coefficient,
+                                 int gaussians,
+                                 double prefactor) {
+  __shared__ double self;
+  if (threadIdx.x == 0) {
+    self = prefactor * contact_value(kernel, coefficient, gaussians);
+  }
+  __syncthreads();
+  return self;
+}
+
+// Per-frame sums over one slice of the owners: the potential of the charged
+// owners in the constrained channel, the number of charged owners and the
+// nine components of sum_i F_i (x) r_i.
+constexpr int kOwnerSums = 11;
+
+__global__ void __launch_bounds__(kFrameThreads)
+    owner_slices(const float* __restrict__ potential,
+                 const float* __restrict__ force,
+                 const double* __restrict__ owner_coord,
+                 const float* __restrict__ owner_charge,
+                 const bool* __restrict__ owner_mask,
+                 const int64_t* __restrict__ owner_frame_ptr,
+                 int channels,
+                 int channel,
+                 int kernel,
+                 const double* __restrict__ coefficient,
+                 int gaussians,
+                 double prefactor,
+                 double* __restrict__ slice_sum) {
+  const double self = self_potential(kernel, coefficient, gaussians, prefactor);
+  const NodeRange slice = slice_of_frame(owner_frame_ptr);
+  double value[kOwnerSums] = {};
+  for (int64_t owner = slice.begin + threadIdx.x; owner < slice.end;
+       owner += kFrameThreads) {
+    if (!owner_mask[owner]) {
+      continue;
+    }
+    if (channel >= 0) {
+      value[0] += double{potential[owner * channels + channel]} -
+                  self * owner_charge[owner * channels + channel];
+    }
+    value[1] += 1.0;
+    for (int a = 0; a < 3; ++a) {
+      for (int b = 0; b < 3; ++b) {
+        value[2 + 3 * a + b] +=
+            double{force[owner * 3 + a]} * owner_coord[owner * 3 + b];
+      }
+    }
+  }
+  write_slice_sum(value, slice_sum);
+}
+
+// Adds the slice sums of owner_slices of every frame in slice order.
+__global__ void owner_totals(const double* __restrict__ slice_sum,
+                             int slice_count,
+                             double* __restrict__ total) {
+  if (threadIdx.x >= kOwnerSums) {
+    return;
+  }
+  double value = 0.0;
+  for (int slice = 0; slice < slice_count; ++slice) {
+    value +=
+        slice_sum[(static_cast(blockIdx.x) * slice_count + slice) *
+                      kOwnerSums +
+                  threadIdx.x];
+  }
+  total[blockIdx.x * kOwnerSums + threadIdx.x] = value;
+}
+
+// Writes the outputs of every owner: the atomic energy q . phi / 2 with phi =
+// phi^rec - prefactor K(0) q, the charge gradient w (phi - [c = ch] mean), and
+// the virial F (x) r plus the share of the rest of the virial W of the frame
+// over its charged nodes. The first threads write the share of the virial of
+// every frame held by the owners, sum_i F_i (x) r_i + n (W - sum_j F_j (x)
+// r_j) / N over the n charged owners of the N charged nodes of the frame.
+// ``owner_local`` holds the owner sums of this domain, ``owner_total`` those
+// of every domain.
+__global__ void finish_owners(const float* __restrict__ potential,
+                              const float* __restrict__ force,
+                              const double* __restrict__ owner_coord,
+                              const float* __restrict__ owner_charge,
+                              const bool* __restrict__ owner_mask,
+                              const int64_t* __restrict__ owner_frame_ptr,
+                              const double* __restrict__ frame_virial,
+                              const double* __restrict__ owner_local,
+                              const double* __restrict__ owner_total,
+                              const double* __restrict__ frame_stat,
+                              int frame_count,
+                              int owner_count,
+                              int channels,
+                              int channel,
+                              int kernel,
+                              const double* __restrict__ coefficient,
+                              int gaussians,
+                              double prefactor,
+                              double* __restrict__ atom_energy,
+                              float* __restrict__ charge_gradient,
+                              float* __restrict__ atom_virial,
+                              float* __restrict__ virial) {
+  const double self = self_potential(kernel, coefficient, gaussians, prefactor);
+  const int start = blockIdx.x * blockDim.x + threadIdx.x;
+  const int threads = gridDim.x * blockDim.x;
+  for (int index = start; index < 9 * frame_count; index += threads) {
+    const int f = index / 9;
+    const int k = index - 9 * f;
+    const double count = frame_stat[f * kStatWidth + kStatCount];
+    const double* local = owner_local + f * kOwnerSums;
+    const double rest =
+        count > 0.0
+            ? (frame_virial[index] - owner_total[f * kOwnerSums + 2 + k]) /
+                  count
+            : 0.0;
+    virial[index] = static_cast(local[2 + k] + local[1] * rest);
+  }
+  for (int owner = start; owner < owner_count; owner += threads) {
+    const int f = frame_of(owner_frame_ptr, frame_count, owner);
+    const bool charged = f >= 0 && owner_mask[owner];
+    // The charged nodes of the frame, of every domain.
+    const double count =
+        charged ? frame_stat[f * kStatWidth + kStatCount] : 1.0;
+    const double mean =
+        charged && channel >= 0 ? owner_total[f * kOwnerSums] / count : 0.0;
+    double energy = 0.0;
+    for (int c = 0; c < channels; ++c) {
+      const int64_t entry = int64_t{owner} * channels + c;
+      const double q = owner_charge[entry];
+      const double phi = charged ? potential[entry] - self * q : 0.0;
+      energy += q * phi;
+      charge_gradient[entry] =
+          static_cast(c == channel ? phi - mean : phi);
+    }
+    atom_energy[owner] = 0.5 * energy;
+    for (int a = 0; a < 3; ++a) {
+      for (int b = 0; b < 3; ++b) {
+        const int k = 3 * a + b;
+        const double moment = double{force[int64_t{owner} * 3 + a]} *
+                              owner_coord[int64_t{owner} * 3 + b];
+        const double rest = charged ? (frame_virial[f * 9 + k] -
+                                       owner_total[f * kOwnerSums + 2 + k]) /
+                                          count
+                                    : 0.0;
+        atom_virial[int64_t{owner} * 9 + k] = static_cast(moment + rest);
+      }
+    }
+  }
+}
+
+IntegerWeights integer_bspline_weights(int order) {
+  IntegerWeights weight{};
+  for (int l = 0; l < order - 1; ++l) {
+    const double x = l + 1.0;
+    double value = 0.0;
+    double binomial = 1.0;
+    for (int j = 0; j <= order; ++j) {
+      if (j > 0) {
+        binomial = binomial * (order - j + 1) / j;
+      }
+      const double base = std::max(x - j, 0.0);
+      value += (j % 2 == 0 ? 1.0 : -1.0) * binomial * std::pow(base, order - 1);
+    }
+    double factorial = 1.0;
+    for (int k = 2; k < order; ++k) {
+      factorial *= k;
+    }
+    weight.value[l] = value / factorial;
+  }
+  return weight;
+}
+
+// The launch of a kernel of float64 work, one item per thread. The float64
+// pipe of a multiprocessor serves every warp on it, so few items spread over
+// as many multiprocessors as there are warps of them, in blocks of a few
+// warps, and many items fill blocks of kThreads.
+struct Float64Launch {
+  int blocks;
+  int threads;
+};
+
+Float64Launch float64_launch(int64_t work) {
+  const int64_t multiprocessors =
+      at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  const int64_t warps = (work + kWarp - 1) / kWarp;
+  const int64_t warps_per_block = std::clamp(
+      (warps + multiprocessors - 1) / multiprocessors, 1, kThreads / kWarp);
+  const int threads = static_cast(warps_per_block * kWarp);
+  return {static_cast(std::max(
+              1, std::min((work + threads - 1) / threads, 65535))),
+          threads};
+}
+
+// Blocks of the influence function of a mesh: a few per multiprocessor fill
+// the device, and no more blocks than the rows of the half spectrum fill, a
+// row to a warp, so that the block sums stay few. The count depends on the
+// mesh and the device alone, so that the modes of a frame are summed in the
+// same order whatever the other frames of its call.
+int influence_blocks(const FrameMesh& mesh) {
+  constexpr int64_t kWarps = kThreads / kWarp;
+  const int64_t rows = int64_t{mesh.n[0]} * mesh.n[1];
+  return static_cast(std::clamp(
+      (rows + kWarps - 1) / kWarps, 1,
+      4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount));
+}
+
+void check_cufft(cufftResult result, const char* operation) {
+  TORCH_CHECK(result == CUFFT_SUCCESS, "latent_charge_sum: cuFFT ", operation,
+              " failed with error ", static_cast(result));
+}
+
+// The real-to-complex and complex-to-real transforms of the meshes of one
+// shape, batched over the channels, and the inverse squared assignment
+// factors of the shape at one B-spline order, which depend on nothing else
+// and are formed once. The plans allocate no work area: every transform runs
+// on a work area of the workspace of its call, set with the stream under
+// ``mutex`` just before the transform is enqueued, so that concurrent calls
+// never share one. A CUDA graph that captures a transform runs it on these
+// plans at every replay, so it holds them while it lives: ``graphs`` counts
+// the graphs that hold them.
+struct MeshPlans {
+  cufftHandle forward = 0;
+  cufftHandle inverse = 0;
+  size_t work_bytes = 0;
+  // float32, n1 + n2 + n3 / 2 + 1 entries on the device (assignment_factor).
+  torch::Tensor factor;
+  c10::DeviceIndex device = 0;
+  std::mutex mutex;
+  std::atomic graphs{0};
+
+  MeshPlans() = default;
+  MeshPlans(const MeshPlans&) = delete;
+  MeshPlans& operator=(const MeshPlans&) = delete;
+  // Plans leave the cache only past kPlanCacheSize shapes and once no graph
+  // holds them; the transforms of eager calls may still be in flight, so the
+  // device finishes them first. The destructor reports no error.
+  ~MeshPlans() {
+    int current = 0;
+    cudaGetDevice(¤t);
+    cudaSetDevice(device);
+    cudaDeviceSynchronize();
+    if (forward != 0) {
+      cufftDestroy(forward);
+    }
+    if (inverse != 0) {
+      cufftDestroy(inverse);
+    }
+    cudaSetDevice(current);
+  }
+};
+
+// The plans of the meshes of ``mesh`` with ``channels`` channels on the
+// current device, with the assignment factors of B-splines of order
+// ``order``, formed on ``stream``, which the creation waits for so that calls
+// on every stream read them.
+std::shared_ptr create_mesh_plans(const FrameMesh& mesh,
+                                             int channels,
+                                             int order,
+                                             cudaStream_t stream) {
+  auto plans = std::make_shared();
+  plans->device = c10::cuda::current_device();
+  long long shape[3] = {mesh.n[0], mesh.n[1], mesh.n[2]};
+  const long long points = shape[0] * shape[1] * shape[2];
+  const long long modes = shape[0] * shape[1] * (shape[2] / 2 + 1);
+  size_t forward_bytes = 0;
+  size_t inverse_bytes = 0;
+  check_cufft(cufftCreate(&plans->forward), "plan creation");
+  check_cufft(cufftSetAutoAllocation(plans->forward, 0), "plan setup");
+  check_cufft(
+      cufftMakePlanMany64(plans->forward, 3, shape, nullptr, 1, points, nullptr,
+                          1, modes, CUFFT_R2C, channels, &forward_bytes),
+      "plan creation");
+  check_cufft(cufftCreate(&plans->inverse), "plan creation");
+  check_cufft(cufftSetAutoAllocation(plans->inverse, 0), "plan setup");
+  check_cufft(
+      cufftMakePlanMany64(plans->inverse, 3, shape, nullptr, 1, modes, nullptr,
+                          1, points, CUFFT_C2R, channels, &inverse_bytes),
+      "plan creation");
+  plans->work_bytes = std::max(forward_bytes, inverse_bytes);
+  const int64_t factor_count =
+      int64_t{mesh.n[0]} + mesh.n[1] + mesh.n[2] / 2 + 1;
+  plans->factor =
+      torch::empty({factor_count}, torch::TensorOptions()
+                                       .dtype(torch::kFloat32)
+                                       .device(torch::kCUDA, plans->device));
+  const Float64Launch launch = float64_launch(factor_count);
+  assignment_factor<<>>(
+      mesh, order, integer_bspline_weights(order),
+      plans->factor.data_ptr());
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  C10_CUDA_CHECK(cudaStreamSynchronize(stream));
+  return plans;
+}
+
+// Ends the hold of a graph on the plans it captured. CUDA runs it on a thread
+// of its own, from which no CUDA call may be made, so the hold is a count
+// that the cache reads rather than an owner that would release the plans.
+void CUDART_CB release_graph_hold(void* plans) {
+  static_cast(plans)->graphs.fetch_sub(1);
+}
+
+// Makes ``graph``, which a capture is recording, hold ``plans`` until it and
+// every graph instantiated from it are destroyed.
+void hold_in_graph(MeshPlans& plans, cudaGraph_t graph) {
+  cudaUserObject_t hold = nullptr;
+  C10_CUDA_CHECK(cudaUserObjectCreate(&hold, &plans, release_graph_hold, 1,
+                                      cudaUserObjectNoDestructorSync));
+  plans.graphs.fetch_add(1);
+  const cudaError_t moved =
+      cudaGraphRetainUserObject(graph, hold, 1, cudaGraphUserObjectMove);
+  if (moved != cudaSuccess) {
+    // The reference stays with this call, whose release ends the hold.
+    cudaUserObjectRelease(hold, 1);
+  }
+  C10_CUDA_CHECK(moved);
+}
+
+// Whether a CUDA graph captures the stream of a call, and that graph.
+struct StreamCapture {
+  cudaStreamCaptureStatus status = cudaStreamCaptureStatusNone;
+  cudaGraph_t graph = nullptr;
+};
+
+StreamCapture stream_capture(cudaStream_t stream) {
+  StreamCapture capture;
+#if CUDART_VERSION >= 12000
+  C10_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &capture.status, nullptr,
+                                          &capture.graph));
+#else
+  C10_CUDA_CHECK(cudaStreamGetCaptureInfo_v2(stream, &capture.status, nullptr,
+                                             &capture.graph));
+#endif
+  return capture;
+}
+
+// The plans of the meshes of ``mesh`` with ``channels`` channels and
+// B-splines of order ``order`` on the current device, kept across calls. A
+// call under ``capture`` makes the graph of the capture hold them and needs
+// them created before, by an eager call, since their creation waits for the
+// stream. The other calls release the least recently used plans that no
+// graph holds past kPlanCacheSize shapes, which a call under capture may not
+// do, since a release waits for the device. The cache itself lives until the
+// process ends and is never destroyed, since its plans would otherwise be
+// released after the CUDA runtime shuts down.
+constexpr size_t kPlanCacheSize = 16;
+
+std::shared_ptr mesh_plans(const FrameMesh& mesh,
+                                      int channels,
+                                      int order,
+                                      const StreamCapture& capture,
+                                      cudaStream_t stream) {
+  using Key = std::array;
+  using Cache = std::list>>;
+  static std::mutex cache_mutex;
+  static Cache& cache = *new Cache();
+  const Key key = {c10::cuda::current_device(),
+                   mesh.n[0],
+                   mesh.n[1],
+                   mesh.n[2],
+                   channels,
+                   order};
+  // Plans that leave the cache are released after the lock, since their
+  // release waits for the device.
+  std::vector> evicted;
+  const std::lock_guard lock(cache_mutex);
+  const auto found =
+      std::find_if(cache.begin(), cache.end(),
+                   [&](const auto& entry) { return entry.first == key; });
+  if (found != cache.end()) {
+    cache.splice(cache.begin(), cache, found);
+  } else {
+    TORCH_CHECK(capture.status == cudaStreamCaptureStatusNone,
+                "latent_charge_sum: a CUDA graph captures only meshes that an "
+                "eager call has evaluated, and none has evaluated the mesh of ",
+                mesh.n[0], " x ", mesh.n[1], " x ", mesh.n[2],
+                " points, channel count ", channels, " and B-spline order ",
+                order);
+    cache.emplace_front(key, create_mesh_plans(mesh, channels, order, stream));
+  }
+  std::shared_ptr plans = cache.front().second;
+  if (capture.status == cudaStreamCaptureStatusActive) {
+    hold_in_graph(*plans, capture.graph);
+  } else if (capture.status == cudaStreamCaptureStatusNone) {
+    // From the least recently used back to, but short of, the front entry,
+    // which holds the plans of this call.
+    auto entry = cache.end();
+    while (cache.size() > kPlanCacheSize && --entry != cache.begin()) {
+      if (entry->second->graphs.load() == 0) {
+        evicted.push_back(std::move(entry->second));
+        entry = cache.erase(entry);
+      }
+    }
+  }
+  return plans;
+}
+
+int blocks_for(int64_t work) {
+  return static_cast(std::max(
+      1, std::min((work + kThreads - 1) / kThreads, 65535)));
+}
+
+// Blocks of a stencil kernel over ``nodes`` nodes, one warp each.
+int stencil_blocks_for(int64_t nodes) {
+  return static_cast(std::max(
+      1, std::min((nodes + kStencilWarps - 1) / kStencilWarps,
+                           int64_t{1} << 20)));
+}
+
+// Calls ``launch`` with the B-spline order as a compile-time constant, for
+// which the stencil kernels unroll their loops.
+template 
+void dispatch_order(int64_t order, const Launch& launch) {
+  switch (order) {
+    case 2:
+      return launch(std::integral_constant{});
+    case 3:
+      return launch(std::integral_constant{});
+    case 4:
+      return launch(std::integral_constant{});
+    case 5:
+      return launch(std::integral_constant{});
+    case 6:
+      return launch(std::integral_constant{});
+    case 7:
+      return launch(std::integral_constant{});
+    case 8:
+      return launch(std::integral_constant{});
+    case 9:
+      return launch(std::integral_constant{});
+    case 10:
+      return launch(std::integral_constant{});
+    case 11:
+      return launch(std::integral_constant{});
+    case 12:
+      return launch(std::integral_constant{});
+    default:
+      TORCH_CHECK(false, "latent_charge_sum: unsupported B-spline order ",
+                  order);
+  }
+}
+
+// The value of ``evaluate`` at the B-spline order as a compile-time constant.
+template 
+auto dispatch_order_value(int64_t order, const Evaluate& evaluate) {
+  decltype(evaluate(std::integral_constant{})) value{};
+  dispatch_order(
+      order, [&](auto order_constant) { value = evaluate(order_constant); });
+  return value;
+}
+
+}  // namespace
+
+namespace deepmd {
+namespace long_range {
+
+ChargeSumOutputs mesh_sum(const torch::Tensor& coord,
+                          const torch::Tensor& charge,
+                          const torch::Tensor& node_mask,
+                          const torch::Tensor& n_node_per_frame,
+                          const torch::Tensor& total_charge,
+                          const torch::Tensor& box,
+                          const torch::Tensor& coefficient,
+                          int64_t kernel,
+                          int64_t channel,
+                          double prefactor,
+                          int64_t order,
+                          double spacing,
+                          int owner_begin,
+                          int owner_count,
+                          const torch::Tensor* communicator) {
+  TORCH_CHECK(order >= 2 && order <= kMaxOrder,
+              "latent_charge_sum: the B-spline order must lie in [2, ",
+              kMaxOrder, "], got ", order);
+  TORCH_CHECK(std::isfinite(spacing) && spacing > 0.0,
+              "latent_charge_sum: the mesh spacing must be positive, got ",
+              spacing);
+  const int64_t node_count = coord.size(0);
+  TORCH_CHECK(node_count < std::numeric_limits::max(),
+              "latent_charge_sum: too many nodes");
+  const int frame_count = static_cast(n_node_per_frame.size(0));
+  const int channels = static_cast(charge.size(1));
+  const int gaussians =
+      kernel == kSog ? static_cast(coefficient.numel() / 2) : 0;
+  const auto double_options = coord.options();
+  const auto float_options = charge.options();
+  const bool every_node = owner_begin == 0 && owner_count == node_count;
+  const auto stream = at::cuda::getCurrentCUDAStream();
+  // The device reads pinned cells in place, and a copy of pageable ones; a
+  // CUDA graph that captures the call reads the cells again at every replay,
+  // which only pinned memory allows.
+  const StreamCapture capture = stream_capture(stream);
+  const bool pinned = box.is_pinned();
+  TORCH_CHECK(capture.status == cudaStreamCaptureStatusNone || pinned,
+              "latent_charge_sum: a CUDA graph that captures periodic frames "
+              "reads their cells at every replay, from box in pinned memory");
+
+  // === Step 1. The meshes of the frames, sized on the host ===
+  std::vector frames(frame_count);
+  std::vector> plans(frame_count);
+  const double* cells = box.data_ptr();
+  int64_t mesh_total = 0;
+  int64_t spectrum_size = 0;
+  int mode_blocks = 1;
+  size_t work_bytes = 0;
+  for (int f = 0; f < frame_count; ++f) {
+    frames[f] = frame_mesh_size(cells + 9 * f, spacing);
+    frames[f].offset = mesh_total;
+    mesh_total += channels * mesh_points(frames[f]);
+    const int half = frames[f].n[2] / 2 + 1;
+    spectrum_size = std::max(spectrum_size, int64_t{channels} * frames[f].n[0] *
+                                                frames[f].n[1] * half);
+    mode_blocks = std::max(mode_blocks, influence_blocks(frames[f]));
+    plans[f] = mesh_plans(frames[f], channels, static_cast(order), capture,
+                          stream);
+    work_bytes = std::max(work_bytes, plans[f]->work_bytes);
+  }
+
+  // === Step 2. The buffers of the call ===
+  const int slice_count = frame_slices(node_count, frame_count);
+  const int owner_slice_count = frame_slices(owner_count, frame_count);
+  const int64_t stencil_bytes = dispatch_order_value(order, [](auto constant) {
+    return static_cast(sizeof(AxisStencil));
+  });
+  Workspace workspace;
+  const int64_t frames_at =
+      workspace.reserve(sizeof(FrameMesh) * frame_count);
+  const int64_t cells_at =
+      workspace.reserve(pinned ? 0 : int64_t{frame_count} * 9);
+  const int64_t frame_ptr_at = workspace.reserve(frame_count + 1);
+  const int64_t slice_charge_at =
+      workspace.reserve(int64_t{frame_count} * slice_count * 2);
+  const int64_t frame_stat_at =
+      workspace.reserve(int64_t{frame_count} * kStatWidth);
+  const int64_t stencils_at =
+      workspace.reserve(3 * node_count * stencil_bytes);
+  const int64_t node_charge_at =
+      workspace.reserve(node_count * channels);
+  const int64_t fixed_mesh_at = workspace.reserve(mesh_total);
+  // The potential mesh of a frame overwrites its charge mesh, which its
+  // forward transform has consumed.
+  const int64_t mesh_at = workspace.reserve(mesh_total);
+  const int64_t spectrum_at = workspace.reserve(spectrum_size);
+  const int64_t work_at = workspace.reserve(work_bytes);
+  const int64_t block_sum_at =
+      workspace.reserve(int64_t{mode_blocks} * kModeSums);
+  const int64_t frame_virial_at =
+      workspace.reserve(int64_t{frame_count} * 9);
+  const int64_t potential_at =
+      workspace.reserve(int64_t{owner_count} * channels);
+  const int64_t owner_slice_at = workspace.reserve(
+      int64_t{frame_count} * owner_slice_count * kOwnerSums);
+  const int64_t owner_local_at =
+      workspace.reserve(int64_t{frame_count} * kOwnerSums);
+  workspace.allocate(coord.device());
+  auto* frames_ptr = workspace.at(frames_at);
+  auto* frame_ptr = workspace.at(frame_ptr_at);
+  auto* slice_charge = workspace.at(slice_charge_at);
+  auto* frame_stat = workspace.at(frame_stat_at);
+  auto* node_charge = workspace.at(node_charge_at);
+  auto* fixed_mesh = workspace.at(fixed_mesh_at);
+  auto* charge_mesh = workspace.at(mesh_at);
+  float* potential_mesh = charge_mesh;
+  auto* spectrum = workspace.at(spectrum_at);
+  void* work = workspace.at(work_at);
+  auto* block_sum = workspace.at(block_sum_at);
+  auto* frame_virial = workspace.at(frame_virial_at);
+  auto* potential = workspace.at(potential_at);
+  auto* owner_slice = workspace.at(owner_slice_at);
+  auto* owner_local = workspace.at(owner_local_at);
+  const double* device_cells = nullptr;
+  if (pinned) {
+    void* mapped = nullptr;
+    C10_CUDA_CHECK(
+        cudaHostGetDevicePointer(&mapped, const_cast(cells), 0));
+    device_cells = static_cast(mapped);
+  } else {
+    // Pageable cells are staged in pinned memory, from which the copy runs
+    // without waiting for the stream; the host allocator keeps the staging
+    // memory until the copy has completed.
+    auto staged =
+        torch::empty({frame_count, 9}, box.options().pinned_memory(true));
+    std::memcpy(staged.data_ptr(), cells,
+                sizeof(double) * 9 * frame_count);
+    auto* copy = workspace.at(cells_at);
+    torch::from_blob(copy, {frame_count, 9}, coord.options())
+        .copy_(staged, /*non_blocking=*/true);
+    device_cells = copy;
+  }
+  for (int f = 0; f < frame_count; ++f) {
+    store_frame_mesh<<<1, kGeometryThreads, 0, stream>>>(
+        frames[f], device_cells + 9 * f, prefactor, frames_ptr + f);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  }
+  C10_CUDA_CHECK(
+      cudaMemsetAsync(fixed_mesh, 0, sizeof(int64_t) * mesh_total, stream));
+
+  // === Step 3. The constraint, and the stencils and charges of the nodes ===
+  const auto counts = n_node_per_frame.to(torch::kInt64).contiguous();
+  frame_offsets<<<1, kOffsetThreads, 0, stream>>>(counts.data_ptr(),
+                                                  frame_count, frame_ptr);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  charge_slices<<>>(
+      charge.data_ptr(), node_mask.data_ptr(), frame_ptr, channels,
+      static_cast(std::max(channel, 0)), slice_charge);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads,
+                 kFrameThreads, 0, stream>>>(
+      slice_charge, slice_count, frame_count, total_charge.data_ptr(),
+      frame_stat);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  dispatch_order(order, [&](auto order_constant) {
+    constexpr int kOrder = decltype(order_constant)::value;
+    auto* node_stencil = workspace.at>(stencils_at);
+    const Float64Launch stencil_launch = float64_launch(3 * node_count);
+    node_stencils
+        <<>>(
+            coord.data_ptr(), charge.data_ptr(),
+            node_mask.data_ptr(), frame_ptr, frames_ptr, frame_stat,
+            frame_count, node_count, channels, static_cast(channel),
+            node_stencil, node_charge);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    // === Step 4. The charge meshes ===
+    spread_charges
+        <<>>(
+            node_mask.data_ptr(), frame_ptr, frames_ptr, node_stencil,
+            node_charge, frame_count, node_count, channels, fixed_mesh);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  });
+  fixed_to_float<<>>(
+      mesh_total, fixed_mesh, charge_mesh);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+  // === Step 5. Influence function and potential meshes, frame by frame ===
+  for (int f = 0; f < frame_count; ++f) {
+    const FrameMesh& grid = frames[f];
+    MeshPlans& plan = *plans[f];
+    {
+      const std::lock_guard lock(plan.mutex);
+      check_cufft(cufftSetStream(plan.forward, stream), "stream setup");
+      check_cufft(cufftSetWorkArea(plan.forward, work), "work area setup");
+      check_cufft(cufftExecR2C(plan.forward, charge_mesh + grid.offset,
+                               reinterpret_cast(spectrum)),
+                  "transform");
+    }
+    const int blocks = influence_blocks(grid);
+    apply_influence<<>>(
+        frames_ptr + f, channels, static_cast(kernel),
+        coefficient.data_ptr(), gaussians,
+        plan.factor.data_ptr(), spectrum, block_sum);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    mode_virial<<<1, kThreads, 0, stream>>>(block_sum, blocks,
+                                            frame_virial + 9 * f);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    {
+      const std::lock_guard lock(plan.mutex);
+      check_cufft(cufftSetStream(plan.inverse, stream), "stream setup");
+      check_cufft(cufftSetWorkArea(plan.inverse, work), "work area setup");
+      check_cufft(
+          cufftExecC2R(plan.inverse, reinterpret_cast(spectrum),
+                       potential_mesh + grid.offset),
+          "transform");
+    }
+  }
+
+  // === Step 6. Potential, force and virial of the owners ===
+  auto force = torch::empty({owner_count, 3}, float_options);
+  if (owner_count > 0) {
+    dispatch_order(order, [&](auto order_constant) {
+      constexpr int kOrder = decltype(order_constant)::value;
+      interpolate_owners<<>>(
+          node_mask.data_ptr(), frame_ptr, frames_ptr,
+          workspace.at>(stencils_at), node_charge,
+          potential_mesh, frame_count, owner_begin, owner_count, channels,
+          potential, force.data_ptr());
+    });
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  }
+  torch::Tensor owner_frame_storage;
+  const int64_t* owner_frame_ptr =
+      owner_frame_offsets(every_node, frame_ptr, owner_count, counts.options(),
+                          owner_frame_storage);
+  const bool* owner_mask = node_mask.data_ptr() + owner_begin;
+  const double* owner_coord = coord.data_ptr() + 3 * owner_begin;
+  const float* owner_charge = node_charge + int64_t{owner_begin} * channels;
+  owner_slices<<>>(potential, force.data_ptr(), owner_coord,
+                           owner_charge, owner_mask, owner_frame_ptr, channels,
+                           static_cast(channel), static_cast(kernel),
+                           coefficient.data_ptr(), gaussians, prefactor,
+                           owner_slice);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  owner_totals<<>>(owner_slice, owner_slice_count,
+                                               owner_local);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  // The mean potential of the constraint and the virial share run over the
+  // charged nodes of every domain.
+  torch::Tensor owner_sum;
+  const double* owner_total = owner_local;
+  if (communicator != nullptr) {
+    owner_sum = sum_over_domains(
+        *communicator, torch::from_blob(owner_local, {frame_count * kOwnerSums},
+                                        double_options));
+    owner_total = owner_sum.data_ptr();
+  }
+  auto atom_energy = torch::empty({owner_count, 1}, double_options);
+  auto charge_gradient = torch::empty({owner_count, channels}, float_options);
+  auto atom_virial = torch::empty({owner_count, 9}, float_options);
+  auto virial = torch::empty({frame_count, 9}, float_options);
+  const Float64Launch owner_launch =
+      float64_launch(std::max(owner_count, 9 * frame_count));
+  finish_owners<<>>(
+      potential, force.data_ptr(), owner_coord, owner_charge, owner_mask,
+      owner_frame_ptr, frame_virial, owner_local, owner_total, frame_stat,
+      frame_count, owner_count, channels, static_cast(channel),
+      static_cast(kernel), coefficient.data_ptr(), gaussians,
+      prefactor, atom_energy.data_ptr(),
+      charge_gradient.data_ptr(), atom_virial.data_ptr(),
+      virial.data_ptr());
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return {atom_energy, charge_gradient, force, atom_virial, virial};
+}
+
+}  // namespace long_range
+}  // namespace deepmd
diff --git a/source/op/pt/long_range/ops.cc b/source/op/pt/long_range/ops.cc
new file mode 100644
index 0000000000..236e61c147
--- /dev/null
+++ b/source/op/pt/long_range/ops.cc
@@ -0,0 +1,99 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+//
+// Operator schema of the latent-charge long-range energy.
+//
+// The schema is declared here, unconditionally, while the CUDA kernel
+// registers itself only against a CUDA-enabled PyTorch (latent_charge_sum.cu),
+// so that the Python front end can tell "library not loaded" from "this device
+// has no kernel".
+//
+// For the nodes of the frames, the operator takes the raw charges c_i of the
+// charge network in C float32 channels, the nodes that carry them (w_i = 1,
+// and 0 elsewhere) and float64 positions r_i. The interactions carry the
+// charges q_i = w_i c_i; when `channel` is a channel index rather than -1,
+// that channel is shifted uniformly over the charged nodes of every frame so
+// that it sums to the total charge Q_f of the frame,
+//
+//   q_i^ch = w_i (c_i^ch - s_f),  s_f = (sum_{j in f} w_j c_j^ch - Q_f) /
+//   max(n_f, 1),
+//
+// with n_f charged nodes in f. Without `box`, every frame is a complete
+// non-periodic system, in which the potential of the charges is the direct
+// sum
+//
+//   phi_i = prefactor * sum_{j != i} q_j K(r_ij);
+//
+// with `box`, whose row f holds the row-major cell of frame f, the rows of the
+// cell being its lattice vectors, phi_i sums over the images of the cell, by
+// the smooth particle-mesh Ewald method (mesh_sum.cu), less the potential
+// prefactor * K(0) q_i of every charge at itself. An all-zero row marks a
+// non-periodic frame, which a call does not mix with periodic ones. Either way
+// the operator returns
+//
+//   atom_energy_i     = q_i . phi_i / 2                        (float64)
+//   charge_gradient_i = dE / dc_i = w_i (phi_i - [c = ch] m_f), m_f the
+//                       mean of phi^ch over the charged nodes of f
+//   force_i           = -dE / dr_i at fixed charges
+//   atom_virial_i     (9 components, row-major)
+//   virial_f          = sum_{i in f} atom_virial_i
+//
+// in float32 unless noted. In a non-periodic frame, with d_ij = r_i - r_j,
+//
+//   force_i  = -prefactor * sum_{j != i} (q_i . q_j) K'(r_ij) / r_ij * d_ij,
+//   W_i      = -prefactor / 2 * sum_{j != i} (q_i . q_j) K'(r_ij) / r_ij
+//              * d_ij (x) d_ij;
+//
+// in a periodic frame the virial of the frame is that of the strain of the
+// cell, W_f = -dE / d(strain), and W_i = F_i (x) r_i + (W_f - sum_{j in f}
+// F_j (x) r_j) / n_f on the charged nodes, the force index first, so that the
+// charged nodes of the frame sum to W_f.
+//
+// The frames occupy the leading sum(n_node_per_frame) <= N nodes, which the
+// caller ensures; the nodes past them are padding and join none. The frame
+// energy is the sum of the atomic energies.
+//
+// `kernel` selects K: 0 is erf(alpha r) / r with `coefficient` holding alpha;
+// 1 is a sum of Gaussians a_k 2^(-c_k r^2), c_k = log2(e) / (2 sigma_k^2),
+// with `coefficient` holding the amplitudes a_k and then the rates c_k.
+// `coefficient` is float64. `total_charge` holds Q_f in float64 and is read
+// only with a constrained channel. `box` is a float64 CPU tensor of shape
+// (F, 9). The mesh of a periodic frame holds along every lattice vector a_d
+// the smallest even number of points without a prime factor above 7 of at
+// least ceil(|a_d| / spacing), and spreads the charges with B-splines of
+// order `order`. The operator sizes the meshes from `box` on the host; the
+// device reads pinned cells in place, while the stream runs the call, and a
+// copy of pageable ones. A CUDA graph that captures a call reads the cells
+// again at every replay, so it needs `box` in pinned memory and evaluates the
+// cells there at replay, provided their meshes keep the sizes of the capture.
+// It captures only meshes whose size, with the channel count and the order,
+// an eager call has evaluated: an eager call forms the cuFFT plans and the
+// assignment factors of a mesh size, which the operator keeps for the sizes
+// used last and for as long as a graph holds them. The caller keeps pinned
+// cells alive and unchanged until the stream has passed a call that reads
+// them. The force and the virial hold the charges fixed; the response of the
+// charges is the caller's, through the charge gradient.
+//
+// `communicator` is an int64 CPU tensor of one element holding the address of
+// the caller's MPI_Comm handle, or zero. A handle of several ranks decomposes
+// one frame into domains, one per rank: the charged nodes of every rank are
+// the charged nodes of the frame, and the sums run over all of them. The
+// outputs of a rank are those of its charged nodes, and W_f is their share of
+// the virial of the frame; the other nodes of the rank carry no charge.
+
+#include 
+
+TORCH_LIBRARY_FRAGMENT(deepmd, library) {
+  library.def(
+      "latent_charge_sum(Tensor coord, Tensor charge, Tensor node_mask, "
+      "Tensor n_node_per_frame, Tensor total_charge, Tensor? box, "
+      "Tensor coefficient, Tensor communicator, int kernel, int channel, "
+      "float prefactor, int order, float spacing) -> "
+      "(Tensor atom_energy, Tensor charge_gradient, Tensor force, "
+      "Tensor atom_virial, Tensor virial)");
+}
+
+// The operator serves inference, whose caller differentiates nothing through
+// it; its outputs carry no gradient.
+TORCH_LIBRARY_IMPL(deepmd, Autograd, library) {
+  library.impl("latent_charge_sum", torch::CppFunction::makeFallthrough());
+}
diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc
index dadd3aab68..c65503da20 100644
--- a/source/op/pt/tabulate_multi_device.cc
+++ b/source/op/pt/tabulate_multi_device.cc
@@ -20,6 +20,43 @@ void GetTensorDevice(const torch::Tensor& t, std::string& str) {
   }
 }
 
+static void ValidateGpuLastLayerSize(const int64_t last_layer_size) {
+  // The GPU tabulation kernels use this dimension either as the block size or
+  // to size dynamic shared memory, so reject invalid models before any launch:
+  // an out-of-range value overflows the block dimension (or divides by zero)
+  // inside the launch configuration.
+  TORCH_CHECK(last_layer_size > 0 && last_layer_size <= 1024,
+              "last_layer_size must be between 1 and 1024 for GPU tabulation, "
+              "but got ",
+              last_layer_size);
+}
+
+static void ValidatePositiveLastLayerSize(const int64_t last_layer_size) {
+  // SE-T-TEBD uses a fixed-size GPU block and loops over the output width, so
+  // it only requires a positive loop bound rather than the generic 1024 cap.
+  TORCH_CHECK(last_layer_size > 0,
+              "last_layer_size must be positive for GPU tabulation, but got ",
+              last_layer_size);
+}
+
+static void ValidateGpuLastLayerSizeBeforeAllocation(
+    const torch::Tensor& table_tensor, const int64_t last_layer_size) {
+  // Forward wrappers allocate descriptor tensors using last_layer_size, so a
+  // negative GPU width must be rejected before torch::empty sees the shape.
+  if (table_tensor.device().is_cuda()) {
+    ValidateGpuLastLayerSize(last_layer_size);
+  }
+}
+
+static void ValidatePositiveGpuLastLayerSizeBeforeAllocation(
+    const torch::Tensor& table_tensor, const int64_t last_layer_size) {
+  // Reject invalid SE-T-TEBD widths before constructing the output tensor,
+  // while preserving widths above the maximum CUDA block dimension.
+  if (table_tensor.device().is_cuda()) {
+    ValidatePositiveLastLayerSize(last_layer_size);
+  }
+}
+
 using NamedTensor = std::pair;
 
 /**
@@ -102,6 +139,7 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor,
   CheckTabulateFusionSeABasisDimension(ndescrpt);
   // compute
   if (device == "GPU") {
+    ValidateGpuLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_a_gpu(descriptor, table, table_info, em_x, em,
                                      two_embed, nloc, nnei, last_layer_size,
@@ -168,6 +206,7 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor,
   CheckTabulateFusionSeABasisDimension(ndescrpt);
   // compute
   if (device == "GPU") {
+    ValidateGpuLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_a_grad_gpu(
         dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy,
@@ -239,6 +278,7 @@ void TabulateFusionSeAGradGradForward(const torch::Tensor& table_tensor,
   CheckTabulateFusionSeABasisDimension(ndescrpt);
   // compute
   if (device == "GPU") {
+    ValidateGpuLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_a_grad_grad_gpu(
         dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
@@ -248,9 +288,6 @@ void TabulateFusionSeAGradGradForward(const torch::Tensor& table_tensor,
         "The input tensor is on the GPU, but the GPU support for the "
         "customized OP library is not enabled.");
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    TORCH_CHECK(last_layer_size <= 1024,
-                "In the process of model compression, the size of the "
-                "last layer of embedding net must be less than 1024!");
   } else if (device == "CPU") {
     deepmd::tabulate_fusion_se_a_grad_grad_cpu(
         dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
@@ -294,6 +331,7 @@ void TabulateFusionSeTForward(const torch::Tensor& table_tensor,
   const int64_t nnei_j = em_tensor.size(2);
   // compute
   if (device == "GPU") {
+    ValidateGpuLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_t_gpu(descriptor, table, table_info, em_x, em,
                                      nloc, nnei_i, nnei_j, last_layer_size);
@@ -346,6 +384,7 @@ void TabulateFusionSeTGradForward(const torch::Tensor& table_tensor,
   const int64_t last_layer_size = descriptor_tensor.size(1);
   // compute
   if (device == "GPU") {
+    ValidateGpuLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_t_grad_gpu(dy_dem_x, dy_dem, table, table_info,
                                           em_x, em, dy, nloc, nnei_i, nnei_j,
@@ -403,6 +442,7 @@ void TabulateFusionSeTGradGradForward(const torch::Tensor& table_tensor,
   const int64_t last_layer_size = descriptor_tensor.size(1);
   // compute
   if (device == "GPU") {
+    ValidateGpuLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_t_grad_grad_gpu(dz_dy, table, table_info, em_x,
                                                em, dz_dy_dem_x, dz_dy_dem, nloc,
@@ -413,9 +453,6 @@ void TabulateFusionSeTGradGradForward(const torch::Tensor& table_tensor,
         "the "
         "customized OP library is not enabled.");
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    TORCH_CHECK(last_layer_size <= 1024,
-                "In the process of model compression, the size of the "
-                "last layer of embedding net must be less than 1024!");
   } else if (device == "CPU") {
     deepmd::tabulate_fusion_se_t_grad_grad_cpu(dz_dy, table, table_info, em_x,
                                                em, dz_dy_dem_x, dz_dy_dem, nloc,
@@ -460,6 +497,7 @@ void TabulateFusionSeTTebdForward(const torch::Tensor& table_tensor,
   const int64_t nnei_j = em_tensor.size(2);
   // compute
   if (device == "GPU") {
+    ValidatePositiveLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_t_tebd_gpu(descriptor, table, table_info, em_x,
                                           em, nloc, nnei_i, nnei_j,
@@ -513,6 +551,7 @@ void TabulateFusionSeTTebdGradForward(const torch::Tensor& table_tensor,
 
   // compute
   if (device == "GPU") {
+    ValidatePositiveLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_t_tebd_grad_gpu(dy_dem_x, table, table_info,
                                                em_x, em, dy, nloc, nnei_i,
@@ -566,6 +605,7 @@ void TabulateFusionSeTTebdGradGradForward(
   const int64_t last_layer_size = descriptor_tensor.size(3);
   // compute
   if (device == "GPU") {
+    ValidatePositiveLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_t_tebd_grad_grad_gpu(
         dz_dy, table, table_info, em_x, em, dz_dy_dem_x, nloc, nnei_i, nnei_j,
@@ -576,11 +616,6 @@ void TabulateFusionSeTTebdGradGradForward(
         "for the "
         "customized OP library is not enabled.");
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    TORCH_CHECK(last_layer_size <= 1024,
-                "In the process of model compression, the "
-                "size of the "
-                "last layer of embedding net must be less "
-                "than 1024!");
   } else if (device == "CPU") {
     deepmd::tabulate_fusion_se_t_tebd_grad_grad_cpu(
         dz_dy, table, table_info, em_x, em, dz_dy_dem_x, nloc, nnei_i, nnei_j,
@@ -616,6 +651,7 @@ void TabulateFusionSeRForward(const torch::Tensor& table_tensor,
   const int64_t nnei = em_tensor.size(1);
   // compute
   if (device == "GPU") {
+    ValidateGpuLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_r_gpu(descriptor, table, table_info, em, nloc,
                                      nnei, last_layer_size);
@@ -662,6 +698,7 @@ void TabulateFusionSeRGradForward(const torch::Tensor& table_tensor,
   const int64_t last_layer_size = descriptor_tensor.size(2);
   // compute
   if (device == "GPU") {
+    ValidateGpuLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_r_grad_gpu(dy_dem, table, table_info, em, dy,
                                           nloc, nnei, last_layer_size);
@@ -708,6 +745,7 @@ void TabulateFusionSeRGradGradForward(const torch::Tensor& table_tensor,
   const int64_t last_layer_size = descriptor_tensor.size(2);
   // compute
   if (device == "GPU") {
+    ValidateGpuLastLayerSize(last_layer_size);
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_r_grad_grad_gpu(
         dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size);
@@ -717,11 +755,6 @@ void TabulateFusionSeRGradGradForward(const torch::Tensor& table_tensor,
         "support for the "
         "customized OP library is not enabled.");
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    TORCH_CHECK(last_layer_size <= 1024,
-                "In the process of model compression, "
-                "the size of the "
-                "last layer of embedding net must be "
-                "less than 1024!");
   } else if (device == "CPU") {
     deepmd::tabulate_fusion_se_r_grad_grad_cpu(
         dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size);
@@ -906,6 +939,7 @@ class TabulateFusionSeAOp
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
       int64_t last_layer_size) {
+    ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())
@@ -1094,6 +1128,7 @@ class TabulateFusionSeAttenOp
       const torch::Tensor& two_embed_tensor,
       int64_t last_layer_size,
       bool is_sorted) {
+    ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())
@@ -1262,6 +1297,7 @@ class TabulateFusionSeTOp
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
       int64_t last_layer_size) {
+    ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())
@@ -1415,6 +1451,7 @@ class TabulateFusionSeROp
       const torch::Tensor& table_info_tensor,
       const torch::Tensor& em_tensor,
       int64_t last_layer_size) {
+    ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())
@@ -1573,6 +1610,8 @@ class TabulateFusionSeTTebdOp
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
       int64_t last_layer_size) {
+    ValidatePositiveGpuLastLayerSizeBeforeAllocation(table_tensor,
+                                                     last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())
diff --git a/source/op/tf/custom_op.h b/source/op/tf/custom_op.h
index 538bfeb2cc..a10b4f4d1b 100644
--- a/source/op/tf/custom_op.h
+++ b/source/op/tf/custom_op.h
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 #pragma once
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -51,6 +52,48 @@ inline Status InvalidArgument(Args&&... args) {
   return tensorflow::errors::InvalidArgument(std::forward(args)...);
 #endif
 }
+
+/**
+ * @brief Derive a dense tensor's per-atom width without truncating division.
+ *
+ * Several low-level TensorFlow ops flatten atom and feature dimensions into a
+ * single axis.  Validate the flattened width before dividing by `nloc`; raw
+ * CPU/GPU kernels cannot safely consume a leftover partial atom row.
+ *
+ * @param per_atom_width Receives the validated feature width for one atom.
+ * @param shape Rank-two tensor shape whose second dimension is flattened.
+ * @param nloc Number of local atoms encoded in the flattened dimension.
+ * @param tensor_name Human-readable input name used in validation errors.
+ * @return An OK status, or InvalidArgument when the width is incompatible.
+ */
+inline Status GetPerAtomWidth(int* per_atom_width,
+                              const TensorShape& shape,
+                              const int nloc,
+                              const char* tensor_name) {
+  const int64_t flattened_width = shape.dim_size(1);
+  if (nloc < 0) {
+    return InvalidArgument("number of local atoms should be non-negative");
+  }
+  if (nloc == 0) {
+    if (flattened_width != 0) {
+      return InvalidArgument(tensor_name,
+                             " width should be zero when nloc is zero");
+    }
+    *per_atom_width = 0;
+    return Status();
+  }
+  if (flattened_width % nloc != 0) {
+    return InvalidArgument(tensor_name, " width ", flattened_width,
+                           " should be divisible by nloc ", nloc);
+  }
+  const int64_t width = flattened_width / nloc;
+  if (width > std::numeric_limits::max()) {
+    return InvalidArgument(tensor_name,
+                           " width per atom exceeds the supported int range");
+  }
+  *per_atom_width = static_cast(width);
+  return Status();
+}
 }  // namespace tf_compat
 }  // namespace deepmd
 
diff --git a/source/op/tf/prod_force_grad_multi_device.cc b/source/op/tf/prod_force_grad_multi_device.cc
index 6678c1b628..39a26cde68 100644
--- a/source/op/tf/prod_force_grad_multi_device.cc
+++ b/source/op/tf/prod_force_grad_multi_device.cc
@@ -73,8 +73,12 @@ class ProdForceSeAGradOp : public OpKernel {
 
     int nframes = net_deriv_tensor.shape().dim_size(0);
     int nloc = natoms(0);
-    int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
-    int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
+    int ndescrpt;
+    int nnei;
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &ndescrpt, net_deriv_shape, nloc, "net deriv"));
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &nnei, nlist_shape, nloc, "nlist"));
 
     // check the sizes
     OP_REQUIRES(
@@ -97,6 +101,9 @@ class ProdForceSeAGradOp : public OpKernel {
     OP_REQUIRES(
         context, (nnei == n_a_sel + n_r_sel),
         deepmd::tf_compat::InvalidArgument("number of neighbors should match"));
+    OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt),
+                deepmd::tf_compat::InvalidArgument(
+                    "descriptor width should be four times neighbor width"));
 
     // Create an output tensor
     TensorShape grad_net_shape;
@@ -194,8 +201,12 @@ class ProdForceSeRGradOp : public OpKernel {
 
     int nframes = net_deriv_tensor.shape().dim_size(0);
     int nloc = natoms(0);
-    int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
-    int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
+    int ndescrpt;
+    int nnei;
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &ndescrpt, net_deriv_shape, nloc, "net deriv"));
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &nnei, nlist_shape, nloc, "nlist"));
 
     // check the sizes
     OP_REQUIRES(
@@ -215,6 +226,9 @@ class ProdForceSeRGradOp : public OpKernel {
                 (int_64(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)),
                 deepmd::tf_compat::InvalidArgument(
                     "number of descriptors should match"));
+    OP_REQUIRES(context, (nnei == ndescrpt),
+                deepmd::tf_compat::InvalidArgument(
+                    "descriptor width should equal neighbor width"));
 
     // Create an output tensor
     TensorShape grad_net_shape;
diff --git a/source/op/tf/prod_force_multi_device.cc b/source/op/tf/prod_force_multi_device.cc
index ae66f13503..b98c6d3593 100644
--- a/source/op/tf/prod_force_multi_device.cc
+++ b/source/op/tf/prod_force_multi_device.cc
@@ -92,8 +92,16 @@ class ProdForceSeAOp : public OpKernel {
     int nloc = natoms[0];
     int nall = natoms[1];
     int nframes = net_deriv_tensor.shape().dim_size(0);
-    int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
-    int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
+    int ndescrpt;
+    int nnei;
+    OP_REQUIRES(context, (nall >= nloc),
+                deepmd::tf_compat::InvalidArgument(
+                    "number of all atoms should be at least nloc"));
+    OP_REQUIRES_OK(context,
+                   deepmd::tf_compat::GetPerAtomWidth(
+                       &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv"));
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &nnei, nlist_tensor.shape(), nloc, "nlist"));
     // check the sizes
     OP_REQUIRES(
         context, (nframes == in_deriv_tensor.shape().dim_size(0)),
@@ -106,6 +114,9 @@ class ProdForceSeAOp : public OpKernel {
         (int_64(nloc) * ndescrpt * 3 == in_deriv_tensor.shape().dim_size(1)),
         deepmd::tf_compat::InvalidArgument(
             "number of descriptors should match"));
+    OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt),
+                deepmd::tf_compat::InvalidArgument(
+                    "descriptor width should be four times neighbor width"));
     // Create an output tensor
     TensorShape force_shape;
     force_shape.AddDim(nframes);
@@ -199,8 +210,16 @@ class ProdForceSeROp : public OpKernel {
     int nloc = natoms[0];
     int nall = natoms[1];
     int nframes = net_deriv_tensor.shape().dim_size(0);
-    int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
-    int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
+    int ndescrpt;
+    int nnei;
+    OP_REQUIRES(context, (nall >= nloc),
+                deepmd::tf_compat::InvalidArgument(
+                    "number of all atoms should be at least nloc"));
+    OP_REQUIRES_OK(context,
+                   deepmd::tf_compat::GetPerAtomWidth(
+                       &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv"));
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &nnei, nlist_tensor.shape(), nloc, "nlist"));
     // check the sizes
     OP_REQUIRES(
         context, (nframes == in_deriv_tensor.shape().dim_size(0)),
@@ -213,6 +232,9 @@ class ProdForceSeROp : public OpKernel {
                  in_deriv_tensor.shape().dim_size(1)),
                 deepmd::tf_compat::InvalidArgument(
                     "number of descriptors should match"));
+    OP_REQUIRES(context, (nnei == ndescrpt),
+                deepmd::tf_compat::InvalidArgument(
+                    "descriptor width should equal neighbor width"));
     // Create an output tensor
     TensorShape force_shape;
     force_shape.AddDim(nframes);
diff --git a/source/op/tf/prod_virial_grad_multi_device.cc b/source/op/tf/prod_virial_grad_multi_device.cc
index fb011cabbd..18c00cc0de 100644
--- a/source/op/tf/prod_virial_grad_multi_device.cc
+++ b/source/op/tf/prod_virial_grad_multi_device.cc
@@ -79,8 +79,12 @@ class ProdVirialSeAGradOp : public OpKernel {
 
     int nframes = net_deriv_tensor.shape().dim_size(0);
     int nloc = natoms(0);
-    int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
-    int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
+    int ndescrpt;
+    int nnei;
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &ndescrpt, net_deriv_shape, nloc, "net deriv"));
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &nnei, nlist_shape, nloc, "nlist"));
 
     // check the sizes
     OP_REQUIRES(
@@ -109,6 +113,9 @@ class ProdVirialSeAGradOp : public OpKernel {
     OP_REQUIRES(
         context, (nnei == n_a_sel + n_r_sel),
         deepmd::tf_compat::InvalidArgument("number of neighbors should match"));
+    OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt),
+                deepmd::tf_compat::InvalidArgument(
+                    "descriptor width should be four times neighbor width"));
 
     // Create an output tensor
     TensorShape grad_net_shape;
@@ -223,8 +230,12 @@ class ProdVirialSeRGradOp : public OpKernel {
 
     int nframes = net_deriv_tensor.shape().dim_size(0);
     int nloc = natoms(0);
-    int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
-    int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
+    int ndescrpt;
+    int nnei;
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &ndescrpt, net_deriv_shape, nloc, "net deriv"));
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &nnei, nlist_shape, nloc, "nlist"));
 
     // check the sizes
     OP_REQUIRES(
@@ -250,6 +261,9 @@ class ProdVirialSeRGradOp : public OpKernel {
     OP_REQUIRES(
         context, (int_64(nloc) * nnei * 3 == rij_shape.dim_size(1)),
         deepmd::tf_compat::InvalidArgument("dim of rij should be  nnei * 3"));
+    OP_REQUIRES(context, (nnei == ndescrpt),
+                deepmd::tf_compat::InvalidArgument(
+                    "descriptor width should equal neighbor width"));
 
     // Create an output tensor
     TensorShape grad_net_shape;
diff --git a/source/op/tf/prod_virial_multi_device.cc b/source/op/tf/prod_virial_multi_device.cc
index 5fff5cac8e..6f1c32628f 100644
--- a/source/op/tf/prod_virial_multi_device.cc
+++ b/source/op/tf/prod_virial_multi_device.cc
@@ -73,9 +73,17 @@ class ProdVirialSeAOp : public OpKernel {
     const int* natoms = natoms_tensor.flat().data();
     int nloc = natoms[0];
     int nall = natoms[1];
-    int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
     int nframes = net_deriv_tensor.shape().dim_size(0);
-    int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
+    int ndescrpt;
+    int nnei;
+    OP_REQUIRES(context, (nall >= nloc),
+                deepmd::tf_compat::InvalidArgument(
+                    "number of all atoms should be at least nloc"));
+    OP_REQUIRES_OK(context,
+                   deepmd::tf_compat::GetPerAtomWidth(
+                       &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv"));
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &nnei, nlist_tensor.shape(), nloc, "nlist"));
     // check the sizes
     OP_REQUIRES(
         context, (nframes == in_deriv_tensor.shape().dim_size(0)),
@@ -94,6 +102,9 @@ class ProdVirialSeAOp : public OpKernel {
     OP_REQUIRES(
         context, (int_64(nloc) * nnei * 3 == rij_tensor.shape().dim_size(1)),
         deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3"));
+    OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt),
+                deepmd::tf_compat::InvalidArgument(
+                    "descriptor width should be four times neighbor width"));
     // Create an output tensor
     TensorShape virial_shape;
     virial_shape.AddDim(nframes);
@@ -179,9 +190,17 @@ class ProdVirialSeROp : public OpKernel {
     const int* natoms = natoms_tensor.flat().data();
     int nloc = natoms[0];
     int nall = natoms[1];
-    int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
     int nframes = net_deriv_tensor.shape().dim_size(0);
-    int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
+    int ndescrpt;
+    int nnei;
+    OP_REQUIRES(context, (nall >= nloc),
+                deepmd::tf_compat::InvalidArgument(
+                    "number of all atoms should be at least nloc"));
+    OP_REQUIRES_OK(context,
+                   deepmd::tf_compat::GetPerAtomWidth(
+                       &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv"));
+    OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth(
+                                &nnei, nlist_tensor.shape(), nloc, "nlist"));
     // check the sizes
     OP_REQUIRES(
         context, (nframes == in_deriv_tensor.shape().dim_size(0)),
@@ -200,6 +219,9 @@ class ProdVirialSeROp : public OpKernel {
     OP_REQUIRES(
         context, (int_64(nloc) * nnei * 3 == rij_tensor.shape().dim_size(1)),
         deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3"));
+    OP_REQUIRES(context, (nnei == ndescrpt),
+                deepmd::tf_compat::InvalidArgument(
+                    "descriptor width should equal neighbor width"));
     // Create an output tensor
     TensorShape virial_shape;
     virial_shape.AddDim(nframes);
diff --git a/source/op/tf/tabulate_multi_device.cc b/source/op/tf/tabulate_multi_device.cc
index 174faf5213..365b63efb7 100644
--- a/source/op/tf/tabulate_multi_device.cc
+++ b/source/op/tf/tabulate_multi_device.cc
@@ -1,5 +1,7 @@
 
 // SPDX-License-Identifier: LGPL-3.0-or-later
+#include 
+
 #include "custom_op.h"
 #include "tabulate.h"
 
@@ -162,6 +164,19 @@ REGISTER_OP("TabulateFusionSeRGradGrad")
     .Input("descriptor: T")
     .Output("dz_dy: T");
 
+template 
+static deepmd::tf_compat::Status validate_gpu_last_layer_size(
+    const int last_layer_size) {
+  // GPU tabulation kernels use this dimension either as the block size or to
+  // size dynamic shared memory, so reject invalid models before any launch.
+  if (std::is_same::value &&
+      (last_layer_size <= 0 || last_layer_size > 1024)) {
+    return deepmd::tf_compat::InvalidArgument(
+        "last_layer_size must be between 1 and 1024 for GPU tabulation");
+  }
+  return deepmd::tf_compat::Status();
+}
+
 template 
 class TabulateFusionSeAOp : public OpKernel {
  public:
@@ -189,6 +204,8 @@ class TabulateFusionSeAOp : public OpKernel {
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 2"));
     OP_REQUIRES(context, (em_tensor.shape().dims() == 3),
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 3"));
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     TensorShape descriptor_shape;
     descriptor_shape.AddDim(em_tensor.shape().dim_size(0));
     descriptor_shape.AddDim(4);  // be careful here;
@@ -248,6 +265,9 @@ class TabulateFusionSeAGradOp : public OpKernel {
     // set size of the sample
     OP_REQUIRES(context, (dy_tensor.shape().dims() == 3),
                 deepmd::tf_compat::InvalidArgument("Dim of table should be 3"));
+    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     int context_output_index = 0;
     Tensor* dy_dem_x_tensor = NULL;
     OP_REQUIRES_OK(context, context->allocate_output(context_output_index++,
@@ -273,7 +293,6 @@ class TabulateFusionSeAGradOp : public OpKernel {
     const FPTYPE* dy = dy_tensor.flat().data();
     const int nloc = em_tensor.shape().dim_size(0);
     const int nnei = em_tensor.shape().dim_size(1);
-    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
 
     if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
@@ -314,6 +333,9 @@ class TabulateFusionSeAGradGradOp : public OpKernel {
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 2"));
     OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 3),
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 3"));
+    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     int context_output_index = 0;
     Tensor* dz_dy_tensor = NULL;
     OP_REQUIRES_OK(context, context->allocate_output(context_output_index++,
@@ -333,7 +355,6 @@ class TabulateFusionSeAGradGradOp : public OpKernel {
     const FPTYPE* dz_dy_dtwo = nullptr;
     const int nloc = em_tensor.shape().dim_size(0);
     const int nnei = em_tensor.shape().dim_size(1);
-    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
 
     if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
@@ -341,10 +362,6 @@ class TabulateFusionSeAGradGradOp : public OpKernel {
           dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
           dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted);
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-      OP_REQUIRES(context, (last_layer_size <= 1024),
-                  deepmd::tf_compat::InvalidArgument(
-                      "In the process of model compression, the size of the "
-                      "last layer of embedding net must be less than 1024!"));
     } else if (device == "CPU") {
       deepmd::tabulate_fusion_se_a_grad_grad_cpu(
           dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
@@ -388,6 +405,8 @@ class TabulateFusionSeAttenOp : public OpKernel {
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 3"));
     OP_REQUIRES(context, (two_embed_tensor.shape().dims() == 2),
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 2"));
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     TensorShape descriptor_shape;
     descriptor_shape.AddDim(em_tensor.shape().dim_size(0));
     descriptor_shape.AddDim(4);  // be careful here;
@@ -453,6 +472,9 @@ class TabulateFusionSeAttenGradOp : public OpKernel {
     // set size of the sample
     OP_REQUIRES(context, (dy_tensor.shape().dims() == 3),
                 deepmd::tf_compat::InvalidArgument("Dim of table should be 3"));
+    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     int context_output_index = 0;
     Tensor* dy_dem_x_tensor = NULL;
     OP_REQUIRES_OK(context, context->allocate_output(context_output_index++,
@@ -482,7 +504,6 @@ class TabulateFusionSeAttenGradOp : public OpKernel {
     const FPTYPE* dy = dy_tensor.flat().data();
     const int nloc = em_tensor.shape().dim_size(0);
     const int nnei = em_tensor.shape().dim_size(1);
-    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
 
     if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
@@ -531,6 +552,9 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel {
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 2"));
     OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 3),
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 3"));
+    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     int context_output_index = 0;
     Tensor* dz_dy_tensor = NULL;
     OP_REQUIRES_OK(context, context->allocate_output(context_output_index++,
@@ -550,7 +574,6 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel {
     const FPTYPE* dz_dy_dtwo = dz_dy_dtwo_tensor.flat().data();
     const int nloc = em_tensor.shape().dim_size(0);
     const int nnei = em_tensor.shape().dim_size(1);
-    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
 
     if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
@@ -558,10 +581,6 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel {
           dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
           dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted);
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-      OP_REQUIRES(context, (last_layer_size <= 1024),
-                  deepmd::tf_compat::InvalidArgument(
-                      "In the process of model compression, the size of the "
-                      "last layer of embedding net must be less than 1024!"));
     } else if (device == "CPU") {
       deepmd::tabulate_fusion_se_a_grad_grad_cpu(
           dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
@@ -603,6 +622,8 @@ class TabulateFusionSeTOp : public OpKernel {
     OP_REQUIRES(
         context, (em_tensor.shape().dims() == 3),
         deepmd::tf_compat::InvalidArgument("Dim of em_tensor should be 3"));
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     TensorShape descriptor_shape;
     descriptor_shape.AddDim(em_tensor.shape().dim_size(0));
     descriptor_shape.AddDim(last_layer_size);
@@ -661,6 +682,9 @@ class TabulateFusionSeTGradOp : public OpKernel {
     OP_REQUIRES(
         context, (dy_tensor.shape().dims() == 2),
         deepmd::tf_compat::InvalidArgument("Dim of dy_tensor should be 2"));
+    const int last_layer_size = descriptor_tensor.shape().dim_size(1);
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     int context_output_index = 0;
     Tensor* dy_dem_x_tensor = NULL;
     OP_REQUIRES_OK(context, context->allocate_output(context_output_index++,
@@ -684,7 +708,6 @@ class TabulateFusionSeTGradOp : public OpKernel {
     const int nloc = em_tensor.shape().dim_size(0);
     const int nnei_i = em_tensor.shape().dim_size(1);
     const int nnei_j = em_tensor.shape().dim_size(2);
-    const int last_layer_size = descriptor_tensor.shape().dim_size(1);
 
     if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
@@ -723,6 +746,9 @@ class TabulateFusionSeTGradGradOp : public OpKernel {
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 2"));
     OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 3),
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 3"));
+    const int last_layer_size = descriptor_tensor.shape().dim_size(1);
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     int context_output_index = 0;
     Tensor* dz_dy_tensor = NULL;
     OP_REQUIRES_OK(context, context->allocate_output(context_output_index++,
@@ -741,7 +767,6 @@ class TabulateFusionSeTGradGradOp : public OpKernel {
     const int nloc = em_tensor.shape().dim_size(0);
     const int nnei_i = em_tensor.shape().dim_size(1);
     const int nnei_j = em_tensor.shape().dim_size(2);
-    const int last_layer_size = descriptor_tensor.shape().dim_size(1);
 
     if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
@@ -749,10 +774,6 @@ class TabulateFusionSeTGradGradOp : public OpKernel {
           dz_dy, table, table_info, em_x, em, dz_dy_dem_x, dz_dy_dem, nloc,
           nnei_i, nnei_j, last_layer_size);
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-      OP_REQUIRES(context, (last_layer_size <= 1024),
-                  deepmd::tf_compat::InvalidArgument(
-                      "In the process of model compression, the size of the "
-                      "last layer of embedding net must be less than 1024!"));
     } else if (device == "CPU") {
       deepmd::tabulate_fusion_se_t_grad_grad_cpu(
           dz_dy, table, table_info, em_x, em, dz_dy_dem_x, dz_dy_dem, nloc,
@@ -787,6 +808,8 @@ class TabulateFusionSeROp : public OpKernel {
                 deepmd::tf_compat::InvalidArgument("Dim of table should be 2"));
     OP_REQUIRES(context, (em_tensor.shape().dims() == 2),
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 2"));
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     TensorShape descriptor_shape;
     descriptor_shape.AddDim(em_tensor.shape().dim_size(0));
     descriptor_shape.AddDim(em_tensor.shape().dim_size(1));  // be careful here;
@@ -842,6 +865,9 @@ class TabulateFusionSeRGradOp : public OpKernel {
     // set size of the sample
     OP_REQUIRES(context, (dy_tensor.shape().dims() == 3),
                 deepmd::tf_compat::InvalidArgument("Dim of table should be 3"));
+    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     int context_output_index = 0;
     Tensor* dy_dem_tensor = NULL;
     OP_REQUIRES_OK(context,
@@ -858,7 +884,6 @@ class TabulateFusionSeRGradOp : public OpKernel {
     const FPTYPE* dy = dy_tensor.flat().data();
     const int nloc = em_tensor.shape().dim_size(0);
     const int nnei = em_tensor.shape().dim_size(1);
-    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
 
     if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
@@ -891,6 +916,9 @@ class TabulateFusionSeRGradGradOp : public OpKernel {
     // set size of the sample
     OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 2),
                 deepmd::tf_compat::InvalidArgument("Dim of input should be 2"));
+    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
+    OP_REQUIRES_OK(context,
+                   validate_gpu_last_layer_size(last_layer_size));
     int context_output_index = 0;
     Tensor* dz_dy_tensor = NULL;
     OP_REQUIRES_OK(context, context->allocate_output(context_output_index++,
@@ -906,17 +934,12 @@ class TabulateFusionSeRGradGradOp : public OpKernel {
     const FPTYPE* dz_dy_dem = dz_dy_dem_tensor.flat().data();
     const int nloc = em_tensor.shape().dim_size(0);
     const int nnei = em_tensor.shape().dim_size(1);
-    const int last_layer_size = descriptor_tensor.shape().dim_size(2);
 
     if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
       deepmd::tabulate_fusion_se_r_grad_grad_gpu(
           dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size);
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-      OP_REQUIRES(context, (last_layer_size <= 1024),
-                  deepmd::tf_compat::InvalidArgument(
-                      "In the process of model compression, the size of the "
-                      "last layer of embedding net must be less than 1024!"));
     } else if (device == "CPU") {
       deepmd::tabulate_fusion_se_r_grad_grad_cpu(
           dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size);
diff --git a/source/tests/common/dpmodel/test_atomic_model_global_stat.py b/source/tests/common/dpmodel/test_atomic_model_global_stat.py
index 0059374e0a..1ddc7d1acc 100644
--- a/source/tests/common/dpmodel/test_atomic_model_global_stat.py
+++ b/source/tests/common/dpmodel/test_atomic_model_global_stat.py
@@ -7,6 +7,9 @@
 from typing import (
     NoReturn,
 )
+from unittest.mock import (
+    patch,
+)
 
 import h5py
 import numpy as np
@@ -330,6 +333,41 @@ def test_preset_bias(self) -> None:
         np.testing.assert_almost_equal(out_bias["foo"], foo_bias)
         np.testing.assert_almost_equal(out_bias["bar"], bar_bias)
 
+        # Standalone change-bias explicitly ignores the checkpoint preset.
+        adjusted_bias = foo_bias + 7.0
+        md0._store_out_stat({"foo": adjusted_bias})
+        forward = md0._get_forward_wrapper_func()
+        seen_biases = []
+
+        def check_forward(*args, **kwargs):
+            bias, _ = md0._fetch_out_stat(["foo"])
+            seen_biases.append(bias["foo"].copy())
+            return forward(*args, **kwargs)
+
+        with patch.object(md0, "_get_forward_wrapper_func", return_value=check_forward):
+            md0.change_out_bias(
+                self.merged_output_stat,
+                bias_adjust_mode="change-by-statistic",
+                preset_out_bias={},
+            )
+        self.assertTrue(seen_biases)
+        for seen_bias in seen_biases:
+            np.testing.assert_almost_equal(seen_bias, adjusted_bias)
+        out_bias, _ = md0._fetch_out_stat(["foo", "bar"])
+        np.testing.assert_almost_equal(
+            out_bias["foo"], np.array([2.0, -5.0]).reshape(2, 1)
+        )
+        assert not np.allclose(out_bias["bar"], bar_bias)
+
+        # The override does not remove the model's configured preset.
+        np.testing.assert_almost_equal(md0.preset_out_bias["foo"], foo_bias)
+        md0.change_out_bias(
+            self.merged_output_stat, bias_adjust_mode="change-by-statistic"
+        )
+        out_bias, _ = md0._fetch_out_stat(["foo", "bar"])
+        np.testing.assert_almost_equal(out_bias["foo"], foo_bias)
+        np.testing.assert_almost_equal(out_bias["bar"], bar_bias)
+
         # 4. a preset that leaves an observed type unassigned is rejected
         md1 = DPDPAtomicModel(
             ds,
diff --git a/source/tests/common/dpmodel/test_data_progress.py b/source/tests/common/dpmodel/test_data_progress.py
new file mode 100644
index 0000000000..797573c264
--- /dev/null
+++ b/source/tests/common/dpmodel/test_data_progress.py
@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Backend-neutral training cursor and selection-RNG checkpoint contract."""
+
+import unittest
+
+from deepmd.dpmodel.train.data_progress import (
+    collect_training_data_progress,
+    preserve_training_data_progress,
+    restore_training_data_progress,
+)
+from deepmd.utils import random as dp_random
+
+
+class _Source:
+    def __init__(self) -> None:
+        self.cursor = 0
+
+    def state_dict(self) -> dict:
+        return {"cursor": self.cursor}
+
+    def load_state_dict(self, state: dict) -> None:
+        self.cursor = state["cursor"]
+
+
+class TestDataProgress(unittest.TestCase):
+    def test_task_cursors_and_selection_rng_roundtrip(self) -> None:
+        sources = {"water": _Source(), "ice": _Source()}
+        sources["water"].cursor = 2
+        sources["ice"].cursor = 5
+        dp_random.seed(17)
+        progress = collect_training_data_progress(sources)
+        expected = [int(dp_random.choice(4)) for _ in range(20)]
+        sources["water"].cursor = 20
+        sources["ice"].cursor = 50
+        restore_training_data_progress(sources, progress)
+        self.assertEqual(sources["water"].cursor, 2)
+        self.assertEqual(sources["ice"].cursor, 5)
+        self.assertEqual([int(dp_random.choice(4)) for _ in range(20)], expected)
+
+    def test_diagnostics_restore_progress_even_on_failure(self) -> None:
+        source = _Source()
+        dp_random.seed(11)
+        before = collect_training_data_progress({"water": source})
+        with self.assertRaisesRegex(RuntimeError, "diagnostic failed"):
+            with preserve_training_data_progress({"water": source}):
+                source.cursor += 1
+                dp_random.choice(10)
+                raise RuntimeError("diagnostic failed")
+        self.assertEqual(collect_training_data_progress({"water": source}), before)
+
+    def test_legacy_and_sources_without_protocol(self) -> None:
+        source = _Source()
+        source.cursor = 3
+        restore_training_data_progress({"water": source}, None)
+        self.assertEqual(source.cursor, 3)
+        sources = {"unsupported": object(), "missing": None}
+        progress = collect_training_data_progress(sources)
+        self.assertEqual(progress["tasks"], {})
+        restore_training_data_progress(sources, progress)
+
+    def test_incompatible_payload_rejected_before_source_restore(self) -> None:
+        sources = {"water": _Source()}
+        progress = collect_training_data_progress(sources)
+        for changed in (
+            {"version": 999},
+            {"tasks": []},
+            {"tasks": {"ice": {"cursor": 5}}},
+        ):
+            with self.subTest(changed=changed), self.assertRaises(ValueError):
+                restore_training_data_progress(sources, progress | changed)
+        del progress["rng"]
+        with self.assertRaises(ValueError):
+            restore_training_data_progress(sources, progress)
+        self.assertEqual(sources["water"].cursor, 0)
diff --git a/source/tests/common/dpmodel/test_fitting_case_embedding.py b/source/tests/common/dpmodel/test_fitting_case_embedding.py
new file mode 100644
index 0000000000..0d7dbc9964
--- /dev/null
+++ b/source/tests/common/dpmodel/test_fitting_case_embedding.py
@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Folding the case embedding of a fitting into its first layer.
+
+The case embedding is the same block of input columns for every atom, so the
+fold moves its product with the first-layer weight into the bias. The folded
+fitting carries no case embedding and gives the outputs of the original, also
+after a serialization round trip. A first layer whose residual adds its input
+to the output keeps the case embedding.
+"""
+
+import numpy as np
+import pytest
+
+from deepmd.dpmodel.fitting import (
+    InvarFitting,
+)
+
+NF, NLOC, NTYPES, NCASE = 2, 5, 2, 3
+
+
+def _fitting(
+    dim_descrpt: int, mixed_types: bool, numb_fparam: int, numb_aparam: int
+) -> InvarFitting:
+    fitting = InvarFitting(
+        "energy",
+        NTYPES,
+        dim_descrpt,
+        1,
+        neuron=[16, 16],
+        resnet_dt=True,
+        numb_fparam=numb_fparam,
+        numb_aparam=numb_aparam,
+        dim_case_embd=NCASE,
+        mixed_types=mixed_types,
+        seed=3,
+    )
+    fitting.set_case_embd(1)
+    return fitting
+
+
+@pytest.mark.parametrize(
+    ("dim_descrpt", "mixed_types", "numb_fparam", "numb_aparam"),
+    [
+        (12, True, 0, 0),  # one network over the descriptor and the case
+        (12, False, 2, 3),  # per-type networks behind frame and atomic parameters
+        (16, True, 0, 0),  # the first layer turns square once folded
+    ],
+)
+def test_the_folded_fitting_gives_the_outputs_of_the_original(
+    dim_descrpt: int, mixed_types: bool, numb_fparam: int, numb_aparam: int
+) -> None:
+    rng = np.random.default_rng(0)
+    fitting = _fitting(dim_descrpt, mixed_types, numb_fparam, numb_aparam)
+    descriptor = rng.normal(size=(NF, NLOC, dim_descrpt))
+    atype = rng.integers(0, NTYPES, size=(NF, NLOC))
+    fparam = rng.normal(size=(NF, numb_fparam)) if numb_fparam else None
+    aparam = rng.normal(size=(NF, NLOC, numb_aparam)) if numb_aparam else None
+    expected = fitting(descriptor, atype, fparam=fparam, aparam=aparam)["energy"]
+
+    fitting.fold_case_embedding()
+
+    assert fitting.dim_case_embd == 0
+    assert fitting.case_embd is None
+    assert fitting.nets[0].in_dim == dim_descrpt + numb_fparam + numb_aparam
+    restored = InvarFitting.deserialize(fitting.serialize())
+    for folded in (fitting, restored):
+        np.testing.assert_allclose(
+            folded(descriptor, atype, fparam=fparam, aparam=aparam)["energy"],
+            expected,
+            rtol=0.0,
+            atol=1e-12,
+        )
+
+
+@pytest.mark.parametrize(
+    "dim_input",
+    [16, 8],  # the residual adds the input once, or twice side by side
+)
+def test_a_residual_over_the_case_columns_keeps_the_case_embedding(
+    dim_input: int,
+) -> None:
+    # The descriptor and the case together are as wide as the first layer, or
+    # half as wide, so its residual adds the case embedding to the output.
+    fitting = _fitting(dim_input - NCASE, True, 0, 0)
+
+    fitting.fold_case_embedding()
+
+    assert fitting.dim_case_embd == NCASE
+    assert fitting.nets[0][0].w.shape == (dim_input, 16)
diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py
index 5630fd5c61..2733c05346 100644
--- a/source/tests/common/dpmodel/test_linear_atomic_model.py
+++ b/source/tests/common/dpmodel/test_linear_atomic_model.py
@@ -47,6 +47,9 @@ def get_dim_fparam(self) -> int:
     def get_dim_aparam(self) -> int:
         return 0
 
+    def has_long_range(self) -> bool:
+        return False
+
     def change_type_map(
         self,
         type_map: list[str],
diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py
index 304a4ba506..27b4c90cdd 100644
--- a/source/tests/common/dpmodel/test_lmdb_data.py
+++ b/source/tests/common/dpmodel/test_lmdb_data.py
@@ -46,9 +46,8 @@
     LmdbTestData,
     LmdbTestDataNlocView,
     _chop_mixed_nloc,
-    _expand_indices_by_blocks,
     _merge_lmdb_chunks,
-    _remap_atom_types,
+    _resample_indices_by_blocks,
     collate_lmdb_frames,
     collect_lmdb_sampling_groups,
     compute_block_targets,
@@ -56,6 +55,7 @@
     decode_lmdb_frame,
     is_lmdb,
     make_neighbor_stat_data,
+    system_block_lookup,
 )
 from deepmd.utils import random as dp_random
 from deepmd.utils.data import (
@@ -1888,6 +1888,124 @@ def setUpClass(cls):
     def tearDownClass(cls):
         cls._tmpdir.cleanup()
 
+    def _subset_type_map_lmdb(
+        self, name: str, second_type: int = 2, *, first_type: int = 2
+    ) -> str:
+        """Build two frames using a sparse source type map and a virtual atom."""
+        path = _create_lmdb_with_virtual_type(f"{self._tmpdir.name}/{name}.lmdb")
+        with lmdb.open(path) as environment, environment.begin(write=True) as txn:
+            meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False)
+            meta.update(
+                nframes=2,
+                type_map=["O", "He", "H"],
+                frame_nlocs=[3, 3],
+                system_info={"natoms": [1, 0, 1]},
+            )
+            txn.put(b"__metadata__", msgpack.packb(meta, use_bin_type=True))
+            frame = msgpack.unpackb(txn.get(b"000000000000"), raw=False)
+            for index, last_type in enumerate((first_type, second_type)):
+                frame["atom_types"] = {
+                    "type": " None:
+        path = self._subset_type_map_lmdb("unused_elements")
+        reader = LmdbDataReader(path, ["H", "O"], batch_size=2)
+        self.addCleanup(reader.close)
+        for index in range(2):
+            np.testing.assert_array_equal(reader[index]["atype"], [1, -1, 0])
+            np.testing.assert_array_equal(reader[index]["natoms"], [3, 3, 1, 1])
+        batch = reader.decode_batch([0, 1])
+        np.testing.assert_array_equal(batch["atype"], [[1, -1, 0], [1, -1, 0]])
+        test_data = LmdbTestData(path, type_map=["H", "O"], shuffle_test=False)
+        np.testing.assert_array_equal(test_data.get_test()["type"], batch["atype"])
+        original_types = LmdbTestData(path, shuffle_test=False)
+        np.testing.assert_array_equal(
+            original_types.get_test()["type"], [[0, -1, 2], [0, -1, 2]]
+        )
+
+    def test_unsupported_element_in_a_later_frame_is_rejected(self) -> None:
+        path = self._subset_type_map_lmdb("used_unsupported_element", second_type=1)
+        reader = LmdbDataReader(path, ["H", "O"], batch_size=2)
+        self.addCleanup(reader.close)
+        np.testing.assert_array_equal(reader[0]["atype"], [1, -1, 0])
+        with self.assertRaisesRegex(ValueError, "frame 1.*He.*model type_map"):
+            reader.__getitem__(1)
+        with self.assertRaisesRegex(ValueError, "frame 1.*He.*model type_map"):
+            reader.decode_batch([0, 1])
+        iterator = LmdbBatchIterator(
+            reader, LmdbBatchSampler(reader, shuffle=False), num_workers=2
+        )
+        self.addCleanup(iterator.close)
+        with self.assertRaisesRegex(ValueError, "frame 1.*He.*model type_map"):
+            next(iterator)
+        test_data = LmdbTestData(path, type_map=["H", "O"], shuffle_test=False)
+        with self.assertRaisesRegex(ValueError, "frame 1.*He.*model type_map"):
+            test_data.get_test()
+        selected = LmdbTestData(
+            path, type_map=["H", "O"], shuffle_test=False, max_frames=1
+        )
+        np.testing.assert_array_equal(selected.get_test()["type"], [[1, -1, 0]])
+
+    def test_batch_selection_ignores_unsupported_first_frame(self) -> None:
+        """Layout discovery reads a frame belonging to the requested batch."""
+        path = self._subset_type_map_lmdb("unsupported_first_batch", first_type=1)
+        for num_workers in (0, 2):
+            with self.subTest(num_workers=num_workers):
+                reader = LmdbDataReader(path, ["H", "O"], batch_size=2)
+                self.addCleanup(reader.close)
+                iterator = LmdbBatchIterator(reader, [[1, 1]], num_workers=num_workers)
+                self.addCleanup(iterator.close)
+                np.testing.assert_array_equal(
+                    next(iterator)["atype"], [[1, -1, 0], [1, -1, 0]]
+                )
+                with self.assertRaisesRegex(ValueError, "frame 0.*He.*model type_map"):
+                    reader.decode_batch([0])
+
+    def test_testdata_probe_ignores_unsupported_first_frame(self) -> None:
+        """PBC discovery is independent of element validation and frame selection."""
+        for box_kind in ("cells", "box", "zero", "missing"):
+            with self.subTest(box_kind=box_kind):
+                path = self._subset_type_map_lmdb(
+                    f"unsupported_first_test_{box_kind}", first_type=1
+                )
+                with (
+                    lmdb.open(path) as environment,
+                    environment.begin(write=True) as txn,
+                ):
+                    frame = msgpack.unpackb(txn.get(b"000000000000"), raw=False)
+                    if box_kind == "box":
+                        frame["box"] = frame.pop("cells")
+                    elif box_kind == "zero":
+                        frame["cells"]["data"] = np.zeros((3, 3)).tobytes()
+                    elif box_kind == "missing":
+                        del frame["cells"]
+                    txn.put(b"000000000000", msgpack.packb(frame, use_bin_type=True))
+                test_data = LmdbTestData(path, type_map=["H", "O"], shuffle_test=False)
+                self.assertEqual(test_data.pbc, box_kind in ("cells", "box"))
+                np.testing.assert_array_equal(
+                    test_data.get_test_by_indices([1])["type"], [[1, -1, 0]]
+                )
+                with self.assertRaisesRegex(ValueError, "frame 0.*He.*model type_map"):
+                    test_data.get_test_by_indices([0])
+
+    def test_subset_remapping_reaches_parallel_decoder_workers(self) -> None:
+        path = self._subset_type_map_lmdb("parallel_subset")
+        reader = LmdbDataReader(path, ["H", "O"], batch_size=2)
+        self.addCleanup(reader.close)
+        iterator = LmdbBatchIterator(
+            reader, LmdbBatchSampler(reader, shuffle=False), num_workers=2
+        )
+        self.addCleanup(iterator.close)
+        np.testing.assert_array_equal(next(iterator)["atype"], [[1, -1, 0], [1, -1, 0]])
+
     def test_reader_no_remap_when_match(self):
         reader = LmdbDataReader(self._lmdb_path, ["O", "H"])
         self.assertIsNone(reader._type_remap)
@@ -1914,8 +2032,10 @@ def test_reader_natoms_vec_after_remap(self):
         self.assertEqual(natoms[3], 2)  # O count
 
     def test_reader_missing_element_raises(self):
+        reader = LmdbDataReader(self._lmdb_path, ["O"])
+        self.addCleanup(reader.close)
         with self.assertRaises(ValueError):
-            LmdbDataReader(self._lmdb_path, ["O"])
+            reader.__getitem__(0)
 
     def test_reader_no_type_map_in_metadata(self):
         tmpdir = tempfile.TemporaryDirectory()
@@ -1945,8 +2065,9 @@ def test_testdata_remap_superset(self):
         self.assertIsNotNone(td._type_remap)
 
     def test_testdata_missing_element_raises(self):
+        td = LmdbTestData(self._lmdb_path, type_map=["O"], shuffle_test=False)
         with self.assertRaises(ValueError):
-            LmdbTestData(self._lmdb_path, type_map=["O"], shuffle_test=False)
+            td.get_test()
 
     def test_testdata_no_type_map_in_metadata(self):
         tmpdir = tempfile.TemporaryDirectory()
@@ -1971,8 +2092,11 @@ def test_virtual_type_preserved_during_remap(self):
 
     def test_positive_out_of_range_type_still_raises(self):
         """A malformed real type must not be mistaken for a virtual sentinel."""
+        path = self._subset_type_map_lmdb("invalid_type_index", second_type=3)
+        reader = LmdbDataReader(path, ["H", "O"])
+        self.addCleanup(reader.close)
         with self.assertRaises(IndexError):
-            _remap_atom_types(np.array([0, 2]), np.array([1, 0]))
+            reader.__getitem__(1)
 
 
 # ============================================================
@@ -1982,7 +2106,7 @@ def test_positive_out_of_range_type_still_raises(self):
 
 class TestAutoProb(unittest.TestCase):
     """Test auto_prob support: frame_system_ids, compute_block_targets,
-    _expand_indices_by_blocks, and LmdbBatchSampler with block_targets.
+    _resample_indices_by_blocks, and LmdbBatchSampler with block_targets.
     """
 
     @classmethod
@@ -2015,6 +2139,48 @@ def test_reader_no_system_ids_backward_compat(self):
         self.assertIsNone(reader.frame_system_ids)
         tmpdir.cleanup()
 
+    def test_uniform_targets_balance_retained_systems(self) -> None:
+        self.assertEqual(
+            compute_block_targets("prob_uniform", 3, [1, 0, 9]),
+            [([0], 5), ([2], 5)],
+        )
+        self.assertEqual(compute_block_targets("prob_uniform", 2, [9, 9]), [])
+        self.assertEqual(compute_block_targets("prob_uniform", 2, [0, 0]), [])
+
+    def test_uniform_target_count_is_bounded(self) -> None:
+        """Uniform quotas stay within one rounding unit per retained system."""
+        counts = [100000] + [1] * 999
+        targets = compute_block_targets("prob_uniform", len(counts), counts)
+        self.assertEqual({target for _, target in targets}, {101})
+        total = sum(target for _, target in targets)
+        self.assertEqual(total, 101000)
+        self.assertGreaterEqual(total, sum(counts))
+        self.assertLess(total, sum(counts) + len(counts))
+
+    def test_uniform_sampler_balances_a_bounded_pass(self) -> None:
+        """A uniform pass can undersample large systems without losing them."""
+        reader = LmdbDataReader(self._lmdb_path, ["O", "H"], batch_size=7)
+        self.addCleanup(reader.close)
+        targets = compute_block_targets(
+            "prob_uniform", reader.nsystems, reader.system_nframes
+        )
+        sampler = LmdbBatchSampler(reader, seed=42, block_targets=targets)
+        seen: set[int] = set()
+        for epoch in range(16):
+            sampler.set_epoch(epoch)
+            expected_batches = len(sampler)
+            batches = list(sampler)
+            self.assertEqual(len(batches), expected_batches)
+            indices = np.concatenate(batches)
+            self.assertEqual(len(indices), 600)
+            np.testing.assert_array_equal(
+                np.bincount(reader.frame_system_ids[indices]), [200, 200, 200]
+            )
+            seen.update(indices.tolist())
+            sampler.set_epoch(epoch)
+            self.assertEqual(list(sampler), batches)
+        self.assertEqual(seen, set(range(len(reader))))
+
     def test_compute_block_targets_equal_weight(self):
         result = compute_block_targets(
             "prob_sys_size;0:1:0.5;1:2:0.5", nsystems=2, system_nframes=[100, 100]
@@ -2067,8 +2233,12 @@ def test_expand_indices_basic(self):
         frame_system_ids = [0] * 5 + [1] * 5
         block_targets = [([0], 25), ([1], 25)]
         rng = np.random.default_rng(42)
-        expanded = _expand_indices_by_blocks(
-            list(range(10)), frame_system_ids, block_targets, rng
+        expanded = _resample_indices_by_blocks(
+            list(range(10)),
+            frame_system_ids,
+            system_block_lookup(block_targets),
+            [target for _, target in block_targets],
+            rng,
         )
         sys0 = [i for i in expanded if frame_system_ids[i] == 0]
         sys1 = [i for i in expanded if frame_system_ids[i] == 1]
@@ -2079,8 +2249,12 @@ def test_expand_indices_no_expansion(self):
         frame_system_ids = [0] * 5 + [1] * 5
         block_targets = [([0], 5), ([1], 5)]
         rng = np.random.default_rng(42)
-        expanded = _expand_indices_by_blocks(
-            list(range(10)), frame_system_ids, block_targets, rng
+        expanded = _resample_indices_by_blocks(
+            list(range(10)),
+            frame_system_ids,
+            system_block_lookup(block_targets),
+            [target for _, target in block_targets],
+            rng,
         )
         self.assertEqual(sorted(expanded), list(range(10)))
 
@@ -2092,8 +2266,12 @@ def test_expand_indices_remainder_sampling(self):
         frame_system_ids = [0] * 10
         block_targets = [([0], 23)]
         rng = np.random.default_rng(42)
-        expanded = _expand_indices_by_blocks(
-            list(range(10)), frame_system_ids, block_targets, rng
+        expanded = _resample_indices_by_blocks(
+            list(range(10)),
+            frame_system_ids,
+            system_block_lookup(block_targets),
+            [target for _, target in block_targets],
+            rng,
         )
         self.assertEqual(len(expanded), 23)
         counts = Counter(expanded)
@@ -2106,8 +2284,12 @@ def test_expand_epoch_diversity(self):
         results = []
         for seed in range(5):
             rng = np.random.default_rng(seed)
-            expanded = _expand_indices_by_blocks(
-                list(range(10)), frame_system_ids, block_targets, rng
+            expanded = _resample_indices_by_blocks(
+                list(range(10)),
+                frame_system_ids,
+                system_block_lookup(block_targets),
+                [target for _, target in block_targets],
+                rng,
             )
             results.append(sorted(expanded[10:]))
         unique = {tuple(r) for r in results}
@@ -2161,6 +2343,16 @@ def get_batch_size_for_nloc(nloc):
         self.assertEqual(sum(map(len, batches)), 3)
         self.assertEqual(sorted(counts), [1, 2])
 
+        sampler = LmdbBatchSampler(
+            TwoSignatureReader(), shuffle=False, seed=42, block_targets=[([0], 1)]
+        )
+        seen: set[int] = set()
+        for _ in range(16):
+            batches = list(sampler)
+            self.assertEqual(sum(map(len, batches)), 1)
+            seen.update(index for batch in batches for index in batch)
+        self.assertEqual(seen, {0, 1})
+
     def test_sampler_tolerates_blocks_that_omit_a_system(self):
         """A system no auto_prob block names is passed through, not rejected.
 
@@ -2875,6 +3067,26 @@ def test_a_second_iterator_does_not_retry_a_lost_pool(self) -> None:
         self._assert_same_batch(batch, self._reader.decode_batch([0, 1, 2, 3]))
         self.assertEqual(pool.submissions, submissions)
 
+    def test_restore_detects_a_stalled_prefetch(self) -> None:
+        """Restoring a cursor must retire a decoder that cannot finish its result."""
+        iterator = LmdbBatchIterator(
+            self._reader, LmdbBatchSampler(self._reader, seed=7), num_workers=2
+        )
+        release = mock.patch.object(lmdb_data_module, "_release_lmdb_executor")
+        release.start()
+        self.addCleanup(release.stop)
+        self.addCleanup(iterator.close)
+        state = iterator.state_dict()
+        entry = lmdb_data_module._LmdbPoolEntry(executor=_StalledPool(), users=1)
+        iterator._pool = entry
+        running = Future()
+        self.assertTrue(running.set_running_or_notify_cancel())
+        iterator._pending = lmdb_data_module._PendingBatch([0, 1], [running])
+        iterator.load_state_dict(state)
+        self.assertFalse(entry.healthy)
+        self.assertIsNone(iterator._pending)
+        self.assertEqual(iterator.state_dict(), state)
+
     @unittest.skipUnless(
         os.name == "posix" and sys.implementation.name == "cpython",
         "requires CPython's POSIX process-pool pipe",
diff --git a/source/tests/common/dpmodel/test_lmdb_data_progress.py b/source/tests/common/dpmodel/test_lmdb_data_progress.py
new file mode 100644
index 0000000000..57a38ca97a
--- /dev/null
+++ b/source/tests/common/dpmodel/test_lmdb_data_progress.py
@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Compact LMDB restart state through the real pt_expt data adapter."""
+
+import copy
+import json
+from contextlib import (
+    ExitStack,
+    closing,
+)
+from unittest.mock import (
+    patch,
+)
+
+import lmdb
+import msgpack
+import pytest
+
+from deepmd.dpmodel.utils import lmdb_data as lmdb_module
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
+
+from .test_lmdb_data import (
+    _create_lmdb,
+    _create_lmdb_with_system_ids,
+    _create_mixed_nloc_lmdb,
+)
+
+
+def _system(path, **kwargs):
+    from deepmd.pt_expt.utils.lmdb_dataset import (
+        LmdbDataSystem,
+    )
+
+    options = {"batch_size": 2, "seed": 7, "num_workers": 0, **kwargs}
+    return LmdbDataSystem(path, ["O", "H"], **options)
+
+
+def _draw(data, count):
+    """Record frame identity and the consumed pass, not the prefetched pass."""
+    return [
+        (data.get_batch()["fid"], data.get_batch_pass_length()) for _ in range(count)
+    ]
+
+
+@pytest.fixture
+def lmdb_path(tmp_path):
+    return _create_lmdb(str(tmp_path / "progress.lmdb"), nframes=10)
+
+
+@pytest.mark.parametrize("workers", [0, 2])
+@pytest.mark.parametrize("seed", [None, 7])
+@pytest.mark.parametrize("consumed", [0, 1, 4, 5, 6])
+def test_logical_cursor_ignores_prefetch(lmdb_path, workers, seed, consumed):
+    """Restart before, inside and across a pass, with real process prefetch."""
+    options = {"num_workers": workers, "seed": seed}
+    with closing(_system(lmdb_path, **options)) as original:
+        _draw(original, consumed)
+        saved = original.state_dict()
+        assert saved["epoch"] == consumed // 5
+        assert saved["batch_index"] == consumed % 5
+        assert isinstance(saved["seed"], int)
+        # Plain bounded metadata, not arrays, decoded data or permutations.
+        assert len(json.dumps(saved)) < 2000
+        expected = _draw(original, 12)
+
+        with closing(_system(lmdb_path, **options)) as resumed:
+            resumed.load_state_dict(saved)
+            assert resumed.state_dict() == saved
+            assert _draw(resumed, 12) == expected
+            if workers and consumed:
+                assert original._batch_iterator.started
+                assert resumed._batch_iterator.started
+
+
+def test_restore_is_lazy_until_trainer_configuration(tmp_path):
+    """Reader/iterator construction predates label requirements and layout."""
+    path = _create_mixed_nloc_lmdb(str(tmp_path / "mixed.lmdb"))
+    with lmdb.open(path) as env, env.begin(write=True) as txn:
+        for index in (1, 4, 7):
+            key = format(index, "012d").encode()
+            frame = msgpack.unpackb(txn.get(key), raw=False)
+            del frame["energies"]
+            txn.put(key, msgpack.packb(frame, use_bin_type=True))
+    requirements = [DataRequirementItem("energy", 1)]
+    options = {"batch_size": "mix:24", "seed": 1}
+    build_batches = lmdb_module._build_all_batches
+    with closing(_system(path, **options)) as original:
+        original.add_data_requirements(requirements)
+        original.use_ragged_batches(True)
+        _draw(original, 3)
+        saved = original.state_dict()
+        expected = _draw(original, 15)
+
+        with (
+            patch.object(
+                lmdb_module, "_build_all_batches", side_effect=AssertionError("sampled")
+            ),
+            closing(_system(path, **options)) as resumed,
+        ):
+            with patch.object(
+                resumed._reader, "decode_batch", side_effect=AssertionError("decoded")
+            ):
+                resumed.load_state_dict(saved)
+                assert resumed.state_dict() == saved
+            # Registration must still be legal, and not plan a sampler pass.
+            resumed.add_data_requirements(requirements)
+            resumed.use_ragged_batches(True)
+            assert resumed._sampler._batches is None
+            # End the build guard before actually asking for the next batch.
+            with patch.object(lmdb_module, "_build_all_batches", new=build_batches):
+                assert _draw(resumed, 15) == expected
+
+
+@pytest.mark.parametrize("workers", [0, 2])
+@pytest.mark.parametrize("consumed", [4, 5, 6, 9])
+def test_ragged_mixed_pass_lengths(tmp_path, workers, consumed):
+    """A five-batch pass followed by a four-batch pass retains its boundary."""
+    path = _create_mixed_nloc_lmdb(str(tmp_path / "mixed.lmdb"))
+    options = {"batch_size": "mix:24", "seed": 1, "num_workers": workers}
+    with closing(_system(path, **options)) as original:
+        original.use_ragged_batches(True)
+        history = _draw(original, consumed)
+        assert history[0][1] == 5
+        if consumed > 5:
+            assert history[5][1] == 4
+        saved = original.state_dict()
+        expected = _draw(original, 15)
+        with closing(_system(path, **options)) as resumed:
+            resumed.load_state_dict(saved)
+            resumed.use_ragged_batches(True)
+            assert resumed.get_batch_pass_length() == expected[0][1]
+            assert _draw(resumed, 15) == expected
+
+
+@pytest.mark.parametrize("seed", [None, 7])
+@pytest.mark.parametrize("consumed", [1, 3, 4])
+def test_equal_rank_cursors_reconstruct_each_shard(lmdb_path, seed, consumed):
+    """Equal progress reconstructs each rank's partition of the same pass."""
+    with ExitStack() as stack:
+        originals = [
+            stack.enter_context(
+                closing(_system(lmdb_path, seed=seed, rank=rank, world_size=2))
+            )
+            for rank in range(2)
+        ]
+        for data in originals:
+            _draw(data, consumed)
+        saved = originals[0].state_dict()
+        assert originals[1].state_dict() == saved
+        for rank, original in enumerate(originals):
+            expected = _draw(original, 10)
+            resumed = stack.enter_context(
+                closing(_system(lmdb_path, seed=seed, rank=rank, world_size=2))
+            )
+            resumed.load_state_dict(saved)
+            assert _draw(resumed, 10) == expected
+
+
+def test_restore_drops_existing_prefetch(lmdb_path):
+    with closing(_system(lmdb_path, num_workers=2)) as data:
+        _draw(data, 3)
+        saved = data.state_dict()
+        expected = _draw(data, 9)
+        pool = data._batch_iterator._pool
+        with (
+            patch.object(lmdb_module, "_build_all_batches") as build,
+            patch.object(data._reader, "decode_batch") as decode,
+            patch.object(lmdb_module, "_release_lmdb_executor") as release,
+        ):
+            data.load_state_dict(saved)
+            build.assert_not_called()
+            decode.assert_not_called()
+            release.assert_not_called()
+        assert data._batch_iterator._pool is pool
+        assert _draw(data, 9) == expected
+
+
+@pytest.mark.parametrize("world_size", [1, 2])
+@pytest.mark.parametrize("diagnostic_reads", [0, 1])
+def test_same_pass_restore_reuses_sampling_plan(
+    lmdb_path, world_size, diagnostic_reads
+):
+    """Display diagnostics must not rebuild a dataset-sized plan each time."""
+    with closing(_system(lmdb_path, batch_size=1, world_size=world_size)) as data:
+        _draw(data, 1)
+        saved = data.state_dict()
+        frame_counts = data._sampler.batch_frame_counts()
+        with patch.object(
+            lmdb_module,
+            "_build_all_batches",
+            side_effect=AssertionError("rebuilt unchanged sampling plan"),
+        ):
+            _draw(data, diagnostic_reads)
+            data.load_state_dict(saved)
+            expected = _draw(data, 1)
+            data.load_state_dict(saved)
+            assert _draw(data, 1) == expected
+            assert data._sampler.batch_frame_counts() is frame_counts
+
+
+@pytest.mark.parametrize("world_size", [1, 2])
+def test_cached_plan_before_layout_selection_is_invalidated(tmp_path, world_size):
+    path = _create_mixed_nloc_lmdb(str(tmp_path / "mixed.lmdb"))
+    options = {"batch_size": "mix:24", "seed": 1, "world_size": world_size}
+    with closing(_system(path, **options)) as original:
+        original.use_ragged_batches(True)
+        _draw(original, 1)
+        saved = original.state_dict()
+        expected = _draw(original, 10)
+        with closing(_system(path, **options)) as resumed:
+            # An earlier metadata query is allowed to have cached the default
+            # rectangular plan before the trainer selects the final layout.
+            assert resumed.nbatches[0] > 0
+            resumed.load_state_dict(saved)
+            resumed.use_ragged_batches(True)
+            assert _draw(resumed, 10) == expected
+
+
+def test_unequal_rank_cursors_restore_independently(lmdb_path):
+    """Local filtering may have consumed a different batch count per rank."""
+    with ExitStack() as stack:
+        saved = []
+        for rank in range(2):
+            original = stack.enter_context(
+                closing(_system(lmdb_path, rank=rank, world_size=2))
+            )
+            _draw(original, rank + 1)
+            saved.append(original.state_dict())
+            expected = _draw(original, 10)
+            resumed = stack.enter_context(
+                closing(_system(lmdb_path, rank=rank, world_size=2))
+            )
+            resumed.load_state_dict(saved[-1])
+            assert _draw(resumed, 10) == expected
+        assert saved[0]["batch_index"] != saved[1]["batch_index"]
+
+
+def test_statistics_do_not_advance_training(lmdb_path):
+    with closing(_system(lmdb_path)) as data:
+        saved = data.state_dict()
+        data.get_stat_batch(0)
+        assert data.state_dict() == saved
+
+
+def test_auto_prob_targets_are_reproduced_and_validated(tmp_path):
+    path = _create_lmdb_with_system_ids(str(tmp_path / "weighted.lmdb"), [4, 8])
+    style = "prob_sys_size;0:1:0.7;1:2:0.3"
+    with closing(_system(path, auto_prob_style=style)) as original:
+        _draw(original, 4)
+        saved = original.state_dict()
+        expected = _draw(original, 12)
+        with closing(_system(path, auto_prob_style=style)) as resumed:
+            resumed.load_state_dict(saved)
+            assert _draw(resumed, 12) == expected
+        with closing(_system(path, auto_prob_style="prob_sys_size")) as changed:
+            with pytest.raises(ValueError, match="block_targets"):
+                changed.load_state_dict(saved)
+
+
+@pytest.mark.parametrize(
+    ("changed", "field"),
+    [
+        ({"batch_size": 3}, "batch_rule"),
+        ({"seed": 12}, "configured_seed"),
+        ({"world_size": 2}, "world_size"),
+    ],
+)
+def test_sampling_configuration_mismatch(lmdb_path, changed, field):
+    with closing(_system(lmdb_path)) as original:
+        _draw(original, 2)
+        saved = original.state_dict()
+        with closing(_system(lmdb_path, **changed)) as resumed:
+            with pytest.raises(ValueError, match=field):
+                resumed.load_state_dict(saved)
+
+
+@pytest.mark.parametrize("field", ["ragged", "availability_keys"])
+def test_late_configuration_mismatch_is_checked_before_decode(tmp_path, field):
+    path = _create_mixed_nloc_lmdb(str(tmp_path / "mixed.lmdb"))
+    with closing(_system(path, batch_size="mix:24")) as original:
+        if field == "ragged":
+            original.use_ragged_batches(True)
+        else:
+            original.add_data_requirements([DataRequirementItem("energy", 1)])
+        _draw(original, 2)
+        saved = original.state_dict()
+        with closing(_system(path, batch_size="mix:24")) as resumed:
+            resumed.load_state_dict(saved)
+            with (
+                patch.object(resumed._reader, "decode_batch") as decode,
+                pytest.raises(ValueError, match=field),
+            ):
+                resumed.get_batch()
+            decode.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("key", "value"),
+    [("version", 99), ("epoch", -1), ("batch_index", 1.5), ("seed", None)],
+)
+def test_invalid_state_is_rejected(lmdb_path, key, value):
+    with closing(_system(lmdb_path)) as data:
+        saved = data.state_dict()
+        invalid = copy.deepcopy(saved)
+        invalid[key] = value
+        with pytest.raises(ValueError):
+            data.load_state_dict(invalid)
+        assert data.state_dict() == saved
+
+
+def test_out_of_range_cursor_is_checked_lazily(lmdb_path):
+    with closing(_system(lmdb_path)) as data:
+        saved = data.state_dict()
+        saved["batch_index"] = 5
+        data.load_state_dict(saved)
+        with pytest.raises(ValueError, match="out of range"):
+            data.get_batch()
+
+
+def test_dataset_transaction_change_is_rejected(lmdb_path):
+    with closing(_system(lmdb_path)) as original:
+        saved = original.state_dict()
+    with lmdb.open(lmdb_path) as env, env.begin(write=True) as txn:
+        key = b"000000000000"
+        frame = msgpack.unpackb(txn.get(key), raw=False)
+        del frame["energies"]
+        txn.put(key, msgpack.packb(frame, use_bin_type=True))
+    with closing(_system(lmdb_path)) as resumed:
+        with pytest.raises(ValueError, match="transaction"):
+            resumed.load_state_dict(saved)
+
+
+def test_different_dataset_is_rejected(lmdb_path, tmp_path):
+    other = _create_lmdb(str(tmp_path / "other.lmdb"), nframes=10)
+    with closing(_system(lmdb_path)) as original, closing(_system(other)) as resumed:
+        with pytest.raises(ValueError, match="path"):
+            resumed.load_state_dict(original.state_dict())
diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py
new file mode 100644
index 0000000000..543ce97dc8
--- /dev/null
+++ b/source/tests/common/dpmodel/test_long_range.py
@@ -0,0 +1,1022 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Primitives of the long-range modules: pair topology, kernels and constraint."""
+
+import itertools
+import pathlib
+
+import numpy as np
+import pytest
+from scipy.integrate import (
+    trapezoid,
+)
+from scipy.special import (
+    erf,
+)
+
+from deepmd.dpmodel.fitting.latent_charge import (
+    LatentChargeFitting,
+)
+from deepmd.dpmodel.long_range import (
+    COULOMB_CONSTANT,
+    BaseInteractionKernel,
+    LatentChargeLongRange,
+    LESKernel,
+    SOGKernel,
+)
+from deepmd.dpmodel.model.model import (
+    get_model,
+)
+from deepmd.dpmodel.utils.exclude_mask import (
+    PairExcludeMask,
+)
+from deepmd.dpmodel.utils.neighbor_graph import (
+    K_BLOCK,
+    build_k_grid,
+    build_pair_graph,
+    build_topologies,
+    enumerate_k_grid_bounds,
+    pair_force_virial,
+    unordered_pair_index,
+)
+from deepmd.utils.argcheck import (
+    long_range_args,
+)
+
+
+def _brute_force_pairs(counts: list[int], mask: np.ndarray | None) -> np.ndarray:
+    pairs = []
+    start = 0
+    for count in counts:
+        nodes = [
+            node for node in range(start, start + count) if mask is None or mask[node]
+        ]
+        pairs += itertools.combinations(nodes, 2)
+        start += count
+    return np.asarray(pairs, dtype=np.int64).reshape(-1, 2).T
+
+
+@pytest.mark.parametrize(
+    "counts", [[0], [1], [2], [3, 0, 1, 5], [7, 200, 1, 0, 33], [1000]]
+)
+@pytest.mark.parametrize("masked", [False, True])
+def test_pair_index_matches_brute_force(counts: list[int], masked: bool) -> None:
+    """Every unordered pair of valid nodes of a frame appears once, in row-major order."""
+    mask = np.random.default_rng(0).random(sum(counts)) > 0.3 if masked else None
+    index = unordered_pair_index(np.asarray(counts), mask)
+    np.testing.assert_array_equal(index, _brute_force_pairs(counts, mask))
+
+
+def test_pair_graph_guards_and_periodic_rejection() -> None:
+    """The pair axis ends in masked guards; a cell is rejected."""
+    coord = np.random.default_rng(1).normal(size=(4, 3))
+    atype = np.array([0, 1, 1, 0])
+    pairs = build_pair_graph(coord, np.array([4]), atype)
+    assert pairs.index.shape == (2, 8)
+    assert pairs.mask.tolist() == [True] * 6 + [False] * 2
+    np.testing.assert_allclose(
+        pairs.vec[:6], coord[pairs.index[0, :6]] - coord[pairs.index[1, :6]]
+    )
+    assert build_topologies((), coord, np.array([4]), atype=atype) == {}
+    with pytest.raises(NotImplementedError, match="non-periodic"):
+        build_topologies(
+            ("all_pairs",), coord, np.array([4]), np.eye(3)[None], atype=atype
+        )
+
+
+def test_pair_force_virial_of_a_central_potential() -> None:
+    """For E = sum 1/r the assembled force is Newtonian and the virial sums."""
+    coord = np.array([[0.0, 0.0, 0.0], [1.0, 0.2, 0.0], [0.0, 1.5, 0.3]])
+    pairs = build_pair_graph(coord, np.array([3]), np.array([0, 1, 0]))
+    # d(1/r)/d(vec) on the real pairs; the zero-length guards carry no gradient.
+    r = np.where(pairs.mask, np.linalg.norm(pairs.vec, axis=-1), 1.0)
+    grad = np.where(pairs.mask[:, None], -pairs.vec / r[:, None] ** 3, 0.0)
+    force, atom_virial, virial = pair_force_virial(
+        grad, pairs.vec, pairs.index, pairs.mask, np.array([3])
+    )
+    expected = np.zeros((3, 3))
+    for i, j in itertools.permutations(range(3), 2):
+        d = coord[i] - coord[j]
+        expected[i] += d / np.linalg.norm(d) ** 3
+    np.testing.assert_allclose(force, expected, rtol=1e-12)
+    np.testing.assert_allclose(atom_virial.sum(0), virial[0], rtol=1e-12)
+    # The virial of central pairs in a cluster is sum_i r_i (x) f_i.
+    np.testing.assert_allclose(virial[0], coord.T @ force, rtol=1e-12, atol=1e-15)
+
+
+def test_pair_force_virial_orients_the_pair_virial() -> None:
+    """The virial of a pair is -g_p (x) pair_vec, the gradient index first."""
+    coord = np.array([[0.0, 0.0, 0.0], [1.0, 0.2, 0.0], [0.0, 1.5, 0.3]])
+    pairs = build_pair_graph(coord, np.array([3]), np.array([0, 1, 0]))
+    # A gradient that is not parallel to its pair, for which the virial is
+    # not symmetric and its orientation shows.
+    rng = np.random.default_rng(1)
+    grad = np.where(pairs.mask[:, None], rng.normal(size=pairs.vec.shape), 0.0)
+    _, atom_virial, virial = pair_force_virial(
+        grad, pairs.vec, pairs.index, pairs.mask, np.array([3])
+    )
+    expected = -np.einsum("pa,pb->ab", grad, pairs.vec)
+    assert not np.allclose(expected, expected.T)
+    np.testing.assert_allclose(virial[0], expected, rtol=1e-12, atol=1e-15)
+    np.testing.assert_allclose(atom_virial.sum(0), expected, rtol=1e-12, atol=1e-15)
+
+
+def test_les_kernel_closed_form() -> None:
+    """LES is erf(alpha r)/r, finite at contact."""
+    kernel = LESKernel(alpha=0.7)
+    r = np.array([0.0, 0.5, 2.0, 9.0])
+    values = kernel(r**2)
+    np.testing.assert_allclose(values[1:], erf(0.7 * r[1:]) / r[1:], rtol=1e-10)
+    np.testing.assert_allclose(values[0], 2 * 0.7 / np.sqrt(np.pi), rtol=1e-10)
+
+
+def test_sog_kernel_closed_form_and_defaults() -> None:
+    """SOG is sum_k a_k exp(-r^2 / 2 sigma_k^2); the defaults discretize 1/r."""
+    kernel = SOGKernel(amplitude=[0.7, -0.2], width=[0.8, 1.7])
+    r = np.array([0.0, 1.1, 3.0])
+    expected = 0.7 * np.exp(-(r**2) / (2 * 0.8**2)) - 0.2 * np.exp(
+        -(r**2) / (2 * 1.7**2)
+    )
+    np.testing.assert_allclose(kernel(r**2), expected, rtol=1e-12)
+    # A scalar amplitude is shared by every Gaussian.
+    shared = SOGKernel(amplitude=0.4, width=[0.8, 1.7])
+    expected = 0.4 * (np.exp(-(r**2) / (2 * 0.8**2)) + np.exp(-(r**2) / (2 * 1.7**2)))
+    np.testing.assert_allclose(shared(r**2), expected, rtol=1e-12)
+    default = SOGKernel(n_gaussian=3, base_width=2.0, width_ratio=1.5)
+    np.testing.assert_allclose(np.exp(default.log_width), [2.0, 3.0, 4.5])
+    np.testing.assert_allclose(
+        default.amplitude, np.sqrt(2 / np.pi) * np.log(1.5) / np.array([2.0, 3.0, 4.5])
+    )
+    # The default kernel exerts the Coulomb force, r^2 K'(r) = -1, between a
+    # few times its first width and a fraction of its widest one.
+    default = SOGKernel()
+    r = np.geomspace(6.5, 150.0, 200)
+    step = 1e-4 * r
+    slope = (default((r + step) ** 2) - default((r - step) ** 2)) / (2 * step)
+    np.testing.assert_allclose(r**2 * slope, -1.0, rtol=0, atol=5e-3)
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"type": "les", "alpha": 0.0},
+        {"type": "sog", "width": [1.0, -2.0]},
+        {"type": "sog", "width": [1.0, 2.0], "amplitude": [1.0, 2.0, 3.0]},
+        {"type": "sog", "width": [1.0, 2.0]},
+        {"type": "sog", "width_ratio": 1.0},
+        {"type": "sog", "amplitude": 1.0, "width_ratio": 0.5},
+    ],
+)
+def test_kernel_parameter_domain(kwargs: dict) -> None:
+    """Widths must be positive and increase by default; amplitudes match the Gaussians."""
+    kernel_cls = BaseInteractionKernel.get_class_by_type(kwargs.pop("type"))
+    with pytest.raises(ValueError):
+        kernel_cls(**kwargs)
+
+
+@pytest.mark.parametrize("kernel", [LESKernel(alpha=1.3), SOGKernel(n_gaussian=4)])
+def test_kernel_round_trip(kernel: BaseInteractionKernel) -> None:
+    """A kernel record restores the same parameters."""
+    restored = BaseInteractionKernel.deserialize(kernel.serialize())
+    r_sq = np.linspace(0.0, 30.0, 7)
+    np.testing.assert_allclose(restored(r_sq), kernel(r_sq), rtol=0, atol=0)
+
+
+def test_kernel_defaults_to_les() -> None:
+    """A configuration without a kernel, or without its type, selects LES."""
+    args = long_range_args()
+    config = args.normalize_value({"type": "latent_charge"}, trim_pattern="_*")
+    assert config["kernel"] == {"type": "les", "alpha": 1.0, "trainable": True}
+    config = args.normalize_value(
+        {"type": "latent_charge", "kernel": {"alpha": 2.0}}, trim_pattern="_*"
+    )
+    assert config["kernel"]["type"] == "les"
+    model = get_model(
+        {
+            "type_map": ["A", "B"],
+            "descriptor": {"type": "dpa4c", "rcut": 3.0, "channels": 8, "seed": 1},
+            "fitting_net": {"type": "ener", "neuron": [8], "seed": 2},
+            "long_range": {"type": "latent_charge", "charge_net": {"neuron": [4]}},
+        }
+    )
+    kernel = model.atomic_model.long_range.kernel
+    assert isinstance(kernel, LESKernel)
+    np.testing.assert_allclose(np.exp(kernel.log_alpha), 1.0)
+
+
+def test_composition_without_the_graph_lower_is_rejected(
+    tmp_path: pathlib.Path,
+) -> None:
+    """A long-range sub-model takes no sibling without the NeighborGraph lower.
+
+    A ``pairtab`` sub-model has no NeighborGraph lower, so the composition
+    would evaluate on the dense route, which a long-range module does not
+    serve; the composition is rejected when it is built.
+    """
+    table = tmp_path / "table.txt"
+    np.savetxt(table, np.column_stack([np.linspace(0.0, 6.0, 61), *[np.zeros(61)] * 3]))
+    child = {
+        "type": "standard",
+        "type_map": ["A", "B"],
+        "descriptor": {"type": "dpa4c", "rcut": 3.0, "channels": 8, "seed": 1},
+        "fitting_net": {"type": "ener", "neuron": [8], "seed": 2},
+        "long_range": {"type": "latent_charge", "charge_net": {"neuron": [4]}},
+    }
+    pairtab = {"type": "pairtab", "tab_file": str(table), "rcut": 5.0, "sel": 20}
+    with pytest.raises(NotImplementedError, match="NeighborGraph"):
+        get_model(
+            {
+                "type": "linear_ener",
+                "type_map": ["A", "B"],
+                "weights": "sum",
+                "models": [child, pairtab],
+            }
+        )
+
+
+def _long_range(constraint: dict | None, channels: int = 2) -> LatentChargeLongRange:
+    charge_net = LatentChargeFitting(
+        "latent_charge", 2, 4, channels, neuron=[4], numb_fparam=1, seed=1
+    )
+    return LatentChargeLongRange(charge_net, LESKernel(), constraint)
+
+
+def test_constraint_shifts_one_channel_over_the_charged_atoms() -> None:
+    """The constrained channel sums to the target over the charged atoms only."""
+    long_range = _long_range({"source": "fparam", "channel": 1, "fparam_index": 0})
+    rng = np.random.default_rng(2)
+    charge = rng.normal(size=(7, 2))
+    node_mask = np.array([True, True, False, True, True, True, False])
+    charge[~node_mask] = 0.0
+    n_node = np.array([4, 3])
+    frame_id = np.array([0, 0, 0, 0, 1, 1, 1])
+    fparam = np.array([[1.5], [-2.0]])
+    constrained = long_range._constrain(
+        charge, node_mask, frame_id, n_node, fparam, None
+    )
+    np.testing.assert_array_equal(constrained[:, 0], charge[:, 0])
+    np.testing.assert_array_equal(constrained[~node_mask], 0.0)
+    totals = np.bincount(frame_id, weights=constrained[:, 1], minlength=2)
+    np.testing.assert_allclose(totals, [1.5, -2.0], rtol=1e-12)
+    # A frame without charged atoms is left as it is.
+    empty = long_range._constrain(
+        charge * 0.0,
+        np.zeros(7, dtype=bool),
+        frame_id,
+        n_node,
+        fparam,
+        None,
+    )
+    np.testing.assert_array_equal(empty, 0.0)
+
+
+@pytest.mark.parametrize(
+    "constraint",
+    [
+        {"source": "total"},
+        {"source": "neutral", "channel": 2},
+        {"source": "fparam", "fparam_index": 1},
+    ],
+)
+def test_constraint_configuration_is_validated(constraint: dict) -> None:
+    """Unknown sources and out-of-range channels or columns are rejected."""
+    with pytest.raises(ValueError):
+        _long_range(constraint)
+
+
+def test_long_range_round_trip() -> None:
+    """The module record holds the charge network, kernel and constraint."""
+    long_range = _long_range({"source": "neutral"})
+    record = long_range.serialize()
+    restored = LatentChargeLongRange.deserialize(record)
+    assert restored.serialize()["charge_constraint"] == record["charge_constraint"]
+    assert restored.required_topologies() == ("all_pairs",)
+    assert not restored.uses_charge_spin()
+
+
+# ---------------------------------------------------------------------------
+# Periodic (k-space) primitives
+# ---------------------------------------------------------------------------
+
+
+def _spherical_ft(
+    kernel: str, k: float, n_r: int = 20000, r_max: float = 60.0
+) -> float:
+    """Isotropic Fourier transform of the kernel by direct radial quadrature."""
+    r = np.linspace(1e-9, r_max, n_r)
+    if kernel == "les":
+        k_r = erf(0.8 * r) / r
+    else:
+        amplitude = np.array([0.7, -0.2, 0.1])
+        width = np.array([0.8, 1.7, 3.2])
+        k_r = np.sum(
+            amplitude[None, :] * np.exp(-(r[:, None] ** 2) / (2 * width**2)), -1
+        )
+    return float(trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(k * r / np.pi), r))
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_kernel_kspace_matches_radial_ft(kernel: str) -> None:
+    """``kernel.kspace`` is the analytic Fourier transform of the real-space kernel."""
+    k_sq = np.array([0.5, 1.0, 2.0, 5.0]) ** 2
+    if kernel == "les":
+        mod = LESKernel(alpha=0.8)
+    else:
+        mod = SOGKernel(amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2])
+    got = mod.kspace(k_sq)
+    for kk, value in zip(np.sqrt(k_sq), got, strict=True):
+        if kernel == "les":
+            # erf(alpha r)/r is not integrable; quadrature of the integrable
+            # complement erfc(alpha r)/r = 1/r - erf(alpha r)/r converges, and
+            # K_hat(k) = 4 pi / k^2 - FT(erfc(alpha r)/r).
+            r = np.linspace(1e-9, 60.0, 20000)
+            k_r = (1.0 - erf(0.8 * r)) / r  # erfc(0.8 r)/r
+            expected = 4.0 * np.pi / kk**2 - float(
+                trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(kk * r / np.pi), r)
+            )
+        else:
+            expected = _spherical_ft(kernel, kk)
+        np.testing.assert_allclose(value, expected, rtol=1e-3, atol=1e-5)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_kernel_contact_value_matches_real_space_limit(kernel: str) -> None:
+    """``kernel.contact_value`` is K(r) at contact."""
+    if kernel == "les":
+        mod = LESKernel(alpha=0.8)
+    else:
+        mod = SOGKernel(amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2])
+    r_sq = np.array([1e-12])
+    np.testing.assert_allclose(
+        mod.contact_value(r_sq), mod(r_sq)[0], rtol=1e-6, atol=1e-9
+    )
+
+
+def test_build_k_grid_orthogonality_and_guards() -> None:
+    """The k points satisfy k . a_d = 2 pi n_d, every frame fills a row of its own."""
+    box = np.array([[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]])[None]
+    grid = build_k_grid(box, 2.0, np.zeros((4, 3)), np.array([4]))
+    # the points fill whole blocks, the last one completed by guard points
+    assert grid.mask.shape == grid.vec.shape[:3]
+    assert grid.vec.shape[2] == K_BLOCK
+    flat_mask = grid.mask.reshape(-1)
+    count = int(flat_mask.sum())
+    assert flat_mask[:count].all() and not flat_mask[count:].any()
+    np.testing.assert_array_equal(grid.slot, [[0, 1, 2, 3]])
+    vec = grid.vec.reshape(-1, 3)[:count]
+    n_integer = vec @ box[0].T / (2 * np.pi)
+    np.testing.assert_allclose(n_integer, np.round(n_integer), atol=1e-10)
+    k_sq = np.sum(vec**2, axis=-1)
+    assert (k_sq > 0).all() and (k_sq <= 4.0 + 1e-12).all()
+    np.testing.assert_allclose(grid.volume, np.abs(np.linalg.det(box[0])))
+    # every integer point inside the sphere is present exactly once
+    ns = np.round(n_integer).astype(int)
+    assert len({tuple(row) for row in ns}) == ns.shape[0]
+    rec = 2 * np.pi * np.linalg.inv(box[0])
+    n_d = np.ceil(2.0 * np.linalg.norm(box[0], axis=1) / (2 * np.pi)).astype(int)
+    count = 0
+    for n0 in range(-n_d[0], n_d[0] + 1):
+        for n1 in range(-n_d[1], n_d[1] + 1):
+            for n2 in range(-n_d[2], n_d[2] + 1):
+                if (n0, n1, n2) != (0, 0, 0):
+                    k = rec @ np.array([n0, n1, n2], dtype=float)
+                    count += int(np.sum(k * k) <= 4.0)
+    assert count == ns.shape[0]
+
+
+def test_build_topologies_periodic_branching() -> None:
+    """A cell switches the pair topology to guards and the k grid on, and back."""
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        ALL_PAIRS,
+        K_GRID,
+    )
+
+    coord = np.random.default_rng(1).normal(size=(4, 3))
+    atype = np.array([0, 1, 1, 0])
+    n_node = np.array([4])
+    box = np.eye(3)[None] * 10.0
+    names = (ALL_PAIRS, K_GRID)
+    params = {K_GRID: {"k_max": 2.0}}
+    with_box = build_topologies(
+        names, coord, n_node, box, atype=atype, topology_params=params
+    )
+    assert with_box[ALL_PAIRS].index.shape[1] == 4  # guards only
+    assert not with_box[ALL_PAIRS].mask.any()
+    assert with_box[K_GRID].mask.sum() > 0
+    without_box = build_topologies(
+        names, coord, n_node, None, atype=atype, topology_params=params
+    )
+    assert without_box[ALL_PAIRS].mask.any()
+    assert not without_box[K_GRID].mask.any()  # guards only
+    # a pair-only declaration still rejects a cell loudly
+    with pytest.raises(NotImplementedError, match="non-periodic"):
+        build_topologies((ALL_PAIRS,), coord, n_node, box, atype=atype)
+    # the k grid needs its build parameter
+    with pytest.raises(ValueError, match="k_max"):
+        build_topologies(names, coord, n_node, box, atype=atype)
+    # the reciprocal-space sum excludes no pair; a non-periodic batch still does
+    pair_excl = PairExcludeMask(2, [(0, 1)])
+    with pytest.raises(NotImplementedError, match="pair_exclude_types"):
+        build_topologies(
+            names,
+            coord,
+            n_node,
+            box,
+            atype=atype,
+            pair_excl=pair_excl,
+            topology_params=params,
+        )
+    excluded = build_topologies(
+        names, coord, n_node, None, atype=atype, pair_excl=pair_excl
+    )
+    assert 0 < excluded[ALL_PAIRS].mask.sum() < without_box[ALL_PAIRS].mask.sum()
+    # a cell without volume is no periodic system
+    with pytest.raises(ValueError, match="no volume"):
+        build_topologies(
+            names, coord, n_node, np.zeros((1, 9)), atype=atype, topology_params=params
+        )
+
+
+def test_build_k_grid_pads_the_rows_of_the_frames() -> None:
+    """Every frame keeps its own points and nodes; shorter rows end in masked guards."""
+    box = np.stack([np.diag([8.0, 7.5, 9.0]), np.diag([5.0, 5.5, 6.0])])
+    n_node = np.array([3, 5])
+    grid = build_k_grid(box, 2.0, np.zeros((8, 3)), n_node)
+    alone = [
+        build_k_grid(box[f : f + 1], 2.0, np.zeros((n, 3)), np.array([n]))
+        for f, n in enumerate(n_node)
+    ]
+    counts = [int(single.mask.sum()) for single in alone]
+    blocks = max(-(-max(counts) // K_BLOCK), 2)
+    assert grid.vec.shape == (2, blocks, K_BLOCK, 3)
+    for f, single in enumerate(alone):
+        mask = grid.mask[f].reshape(-1)
+        vec = grid.vec[f].reshape(-1, 3)
+        np.testing.assert_array_equal(mask, np.arange(blocks * K_BLOCK) < counts[f])
+        np.testing.assert_allclose(
+            vec[: counts[f]], single.vec[0].reshape(-1, 3)[: counts[f]]
+        )
+        assert np.all(vec[counts[f] :] == [1e-6, 0.0, 0.0])
+    np.testing.assert_array_equal(grid.slot, [[0, 1, 2, -1, -1], [3, 4, 5, 6, 7]])
+
+
+@pytest.mark.parametrize("kernel", [LESKernel(alpha=0.6), SOGKernel(n_gaussian=3)])
+def test_kspace_energy_of_padded_frames_matches_frames_alone(
+    kernel: BaseInteractionKernel,
+) -> None:
+    """Frames of unequal node counts, an empty one among them, sum as they do alone."""
+    from deepmd.dpmodel.long_range.reciprocal import (
+        kspace_energy,
+    )
+
+    box = np.stack(
+        [np.diag([8.0, 7.5, 9.0]), np.diag([6.0, 6.5, 7.0]), np.diag([5.0, 5.5, 6.0])]
+    )
+    n_node = np.array([4, 0, 6])
+    rng = np.random.default_rng(3)
+    coord = rng.uniform(0.0, 5.0, (10, 3))
+    charge = rng.normal(0.0, 0.5, (10, 2))
+    energy = kspace_energy(charge, build_k_grid(box, 2.5, coord, n_node), kernel)
+    start = 0
+    for f, n in enumerate(n_node):
+        if n == 0:
+            continue
+        nodes = slice(start, start + n)
+        alone = build_k_grid(box[f : f + 1], 2.5, coord[nodes], np.array([n]))
+        np.testing.assert_allclose(
+            energy[nodes], kspace_energy(charge[nodes], alone, kernel), rtol=1e-12
+        )
+        start += n
+
+
+def test_kspace_energy_keeps_the_self_term_without_wave_numbers() -> None:
+    """A periodic frame without wave numbers below the cutoff keeps its self term.
+
+    For a small cell and a wide kernel, every wave number of the cell lies
+    beyond the cutoff, so the lattice sum reduces to the self term
+    -k_e K(0) q_i^2 / 2 of every node, which the mesh evaluation gives as
+    well, over both the ragged and the fixed k grid; a non-periodic batch has
+    zero energy.
+    """
+    from deepmd.dpmodel.long_range.reciprocal import (
+        kspace_energy,
+        mesh_energy,
+    )
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        build_mesh_grid,
+    )
+    from deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph import (
+        build_fixed_k_grid,
+        empty_k_grid,
+    )
+
+    kernel = LESKernel(alpha=0.03)
+    k_max = kernel.kspace_cutoff()
+    box = np.diag([3.0, 3.0, 3.0])[None]
+    assert 2.0 * np.pi / 3.0 > k_max
+    coord = np.array([[0.2, 0.3, 0.4], [1.7, 1.6, 1.5]])
+    charge = np.array([[1.0], [-1.0]])
+    n_node = np.array([2])
+    expected = -0.5 * COULOMB_CONSTANT * 2.0 * 0.03 / np.sqrt(np.pi) * charge[:, 0] ** 2
+    for grid in (
+        build_k_grid(box, k_max, coord, n_node),
+        build_fixed_k_grid(box, np.array([1, 1, 1]), k_max, coord, n_node),
+    ):
+        assert not np.any(grid.mask)
+        np.testing.assert_allclose(
+            kspace_energy(charge, grid, kernel), expected, rtol=1e-12
+        )
+    mesh = build_mesh_grid(box, 0.5, coord)
+    frame_id = np.zeros(2, dtype=np.int64)
+    np.testing.assert_allclose(
+        mesh_energy(charge, mesh, frame_id, kernel, 8), expected, rtol=1e-4
+    )
+    np.testing.assert_array_equal(
+        kspace_energy(charge, empty_k_grid(1, coord), kernel), 0.0
+    )
+
+
+# ---------------------------------------------------------------------------
+# Fixed k-grid mode
+# ---------------------------------------------------------------------------
+
+
+def _module_topologies(
+    module: LatentChargeLongRange,
+    coord: np.ndarray,
+    n_node: np.ndarray,
+    box: np.ndarray | None,
+    atype: np.ndarray,
+) -> dict:
+    """The topologies of a module, built with its parameters."""
+    return build_topologies(
+        module.required_topologies(),
+        coord,
+        n_node,
+        box,
+        atype=atype,
+        topology_params=module.topology_params(),
+    )
+
+
+def _periodic_long_range(**kwargs: object) -> LatentChargeLongRange:
+    """A periodic module, over the explicit wave numbers unless ``k_grid_mode`` says otherwise."""
+    charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1)
+    kwargs.setdefault("k_grid_mode", "ragged")
+    return LatentChargeLongRange(
+        charge_net, LESKernel(alpha=0.3), None, periodic=True, **kwargs
+    )
+
+
+def _cells() -> np.ndarray:
+    """Two cells of different sizes, the second the larger one."""
+    return np.array(
+        [
+            [[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]],
+            [[12.0, 0.5, 0.0], [0.2, 11.0, 0.3], [0.1, 0.4, 10.5]],
+        ]
+    )
+
+
+def test_enumerate_k_grid_bounds() -> None:
+    """The grid is the lexicographic meshgrid of the bounds, without the zero vector."""
+    n_grid = enumerate_k_grid_bounds([2, 1, 0])
+    assert n_grid.shape == (5 * 3 * 1 - 1, 3)
+    assert not np.any(np.all(n_grid == 0, axis=-1))
+    expected = [
+        (n0, n1, 0) for n0 in range(-2, 3) for n1 in range(-1, 2) if (n0, n1) != (0, 0)
+    ]
+    np.testing.assert_array_equal(n_grid, np.asarray(expected, dtype=np.int64))
+    with pytest.raises(ValueError, match="bounds"):
+        enumerate_k_grid_bounds([1, -1, 0])
+
+
+def test_fixed_k_grid_matches_ragged_enumeration() -> None:
+    """The masked points of the fixed grid are exactly the ragged k points."""
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        build_fixed_k_grid,
+        enumerate_k_grid_bounds,
+        k_grid_bounds_of_cells,
+    )
+
+    box = _cells()
+    k_max = 2.0
+    bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0)
+    nk = enumerate_k_grid_bounds(bounds).shape[0]
+    n_node = np.array([4, 5])
+    fixed = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3)), n_node)
+    # a static layout: the blocks of the grid in every row
+    assert fixed.vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3)
+    ragged = build_k_grid(box, k_max, np.zeros((9, 3)), n_node)
+    for frame in range(2):
+        expected = {
+            tuple(row) for row in np.round(ragged.vec[frame][ragged.mask[frame]], 8)
+        }
+        got = {tuple(row) for row in np.round(fixed.vec[frame][fixed.mask[frame]], 8)}
+        assert got == expected
+    # the smaller cell masks more points of the shared grid
+    assert fixed.mask[0].sum() < fixed.mask[1].sum()
+    np.testing.assert_allclose(fixed.volume, np.abs(np.linalg.det(box)), rtol=1e-12)
+
+
+def test_fixed_k_grid_truncates_oversized_cells() -> None:
+    """A cell beyond the grid is served by the grid alone (truncated k-sum), with a warning."""
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        build_fixed_k_grid,
+        enumerate_k_grid_bounds,
+        k_grid_bounds_of_cells,
+    )
+
+    box = _cells()
+    k_max = 2.0
+    bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0)
+    nk = enumerate_k_grid_bounds(bounds).shape[0]
+    n_node = np.array([4, 5])
+    with pytest.warns(UserWarning, match="truncated"):
+        bigger = build_fixed_k_grid(box * 2.0, bounds, k_max, np.zeros((9, 3)), n_node)
+    # the layout is untouched: static, every grid point present
+    assert bigger.vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3)
+    # every grid point is still a valid wave number of the larger cell; the
+    # mask keeps the whole k_max sphere the grid reaches
+    k2 = np.sum(bigger.vec**2, axis=-1)
+    real = (np.arange(bigger.mask[0].size) < nk).reshape(bigger.mask.shape[1:])
+    np.testing.assert_array_equal(bigger.mask, real & (k2 <= k_max * k_max))
+    # the larger cell activates more of the shared grid than the original did
+    original = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3)), n_node)
+    assert bigger.mask.sum() > original.mask.sum()
+    # a larger wave-number cutoff (a drifting kernel) truncates the same way
+    with pytest.warns(UserWarning, match="truncated"):
+        build_fixed_k_grid(box, bounds, k_max * 2.0, np.zeros((9, 3)), n_node)
+
+
+def test_fixed_k_grid_mode_configuration() -> None:
+    """The fixed mode requires periodicity, and the bounds the fixed mode."""
+    with pytest.raises(ValueError, match="k_grid_mode"):
+        _periodic_long_range(k_grid_mode="dense")
+    with pytest.raises(ValueError, match="periodic"):
+        LatentChargeLongRange(
+            LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1),
+            LESKernel(),
+            None,
+            k_grid_mode="fixed",
+        )
+    with pytest.raises(ValueError, match="k_grid_bounds"):
+        _periodic_long_range(k_grid_bounds=[3, 3, 3])
+    # explicit bounds build the grid at construction, no statistics needed
+    long_range = _periodic_long_range(k_grid_mode="fixed", k_grid_bounds=[2, 2, 2])
+    np.testing.assert_array_equal(long_range.k_grid_bounds, [2, 2, 2])
+    with pytest.raises(ValueError, match="non-negative"):
+        _periodic_long_range(k_grid_mode="fixed", k_grid_bounds=[0, 0, 0])
+
+
+def test_fixed_k_grid_stats_size_the_grid() -> None:
+    """The data statistics size the grid by the largest sampled cell."""
+    long_range = _periodic_long_range(k_grid_mode="fixed")
+    # the grid is not sized before the statistics run
+    coord = np.zeros((9, 3))
+    atype = np.zeros(9, dtype=np.int64)
+    with pytest.raises(ValueError, match="not sized yet"):
+        _module_topologies(long_range, coord, np.array([4, 5]), _cells(), atype)
+    long_range.compute_input_stats(
+        [{"box": _cells()[:1].reshape(1, 9)}, {"box": _cells()}]
+    )
+    k_max = long_range.kernel.kspace_cutoff(long_range.kspace_tolerance)
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        k_grid_bounds_of_cells,
+    )
+
+    expected = np.max(k_grid_bounds_of_cells(_cells(), k_max), axis=0)
+    np.testing.assert_array_equal(long_range.k_grid_bounds, expected)
+    topologies = _module_topologies(
+        long_range, coord, np.array([4, 5]), _cells(), atype
+    )
+    nk = np.prod(2 * expected + 1) - 1
+    assert topologies["k_grid"].vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3)
+    # statistics without a cell cannot size the grid
+    unsized = _periodic_long_range(k_grid_mode="fixed")
+    with pytest.raises(ValueError, match="no simulation cell"):
+        unsized.compute_input_stats([{"coord": np.zeros((1, 3))}])
+
+
+def test_fixed_k_grid_round_trip() -> None:
+    """The record keeps the grid mode and the sized grid."""
+    long_range = _periodic_long_range(k_grid_mode="fixed")
+    long_range.compute_input_stats([{"box": _cells()}])
+    record = long_range.serialize()
+    assert record["k_grid_mode"] == "fixed"
+    restored = LatentChargeLongRange.deserialize(record)
+    assert restored.k_grid_mode == "fixed"
+    np.testing.assert_array_equal(restored.k_grid_bounds, long_range.k_grid_bounds)
+    # the ragged mode serializes no grid bounds
+    plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize())
+    assert plain.k_grid_mode == "ragged"
+    assert plain.k_grid_bounds is None
+    # a periodic module evaluates on the mesh unless told otherwise, and a
+    # periodic record without a mode is one of the explicit wave numbers
+    charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1)
+    default = LatentChargeLongRange(charge_net, LESKernel(), None, periodic=True)
+    assert default.k_grid_mode == "pme"
+    assert default.required_topologies() == ("all_pairs", "pme_mesh")
+    record = default.serialize()
+    assert record["k_grid_mode"] == "pme"
+    record.pop("k_grid_mode")
+    assert LatentChargeLongRange.deserialize(record).k_grid_mode == "ragged"
+
+
+# ---------------------------------------------------------------------------
+# Decoupled k grid
+# ---------------------------------------------------------------------------
+
+
+def test_decoupled_k_grid_requires_periodic() -> None:
+    """Decoupling is only meaningful for a periodic module."""
+    with pytest.raises(ValueError, match="decouple_k_grid"):
+        LatentChargeLongRange(
+            LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1),
+            LESKernel(),
+            None,
+            decouple_k_grid=True,
+        )
+
+
+def test_decoupled_k_grid_ignores_kernel_drift() -> None:
+    """A decoupled grid keeps the k points of the initial kernel however the
+    kernel parameters move; a coupled one re-derives them every batch.
+    """
+    coord = np.random.default_rng(3).normal(size=(5, 3))
+    atype = np.zeros(5, dtype=np.int64)
+    n_node = np.array([5])
+    box = np.eye(3)[None] * 8.0
+
+    decoupled = _periodic_long_range(decouple_k_grid=True)
+    coupled = _periodic_long_range()
+    # the frozen cutoff is the cutoff of the initial kernel
+    assert decoupled.frozen_k_max == pytest.approx(
+        LESKernel(alpha=0.3).kspace_cutoff(1e-8)
+    )
+    initial = _module_topologies(decoupled, coord, n_node, box, atype)["k_grid"]
+    reference = _module_topologies(coupled, coord, n_node, box, atype)["k_grid"]
+    np.testing.assert_array_equal(initial.vec, reference.vec)
+    # a larger alpha moves the cutoff out, growing the coupled grid
+    for module in (decoupled, coupled):
+        module.kernel.log_alpha = np.log(1.5)
+    drifted = _module_topologies(decoupled, coord, n_node, box, atype)["k_grid"]
+    moved = _module_topologies(coupled, coord, n_node, box, atype)["k_grid"]
+    np.testing.assert_array_equal(drifted.vec, initial.vec)
+    np.testing.assert_array_equal(drifted.mask, initial.mask)
+    assert np.sum(moved.mask) > np.sum(initial.mask)
+
+
+def test_decoupled_fixed_grid_keeps_its_mask() -> None:
+    """The mask of a decoupled fixed grid ignores the kernel drift."""
+    coord = np.zeros((5, 3))
+    atype = np.zeros(5, dtype=np.int64)
+    n_node = np.array([5])
+    box = np.eye(3)[None] * 8.0
+
+    long_range = _periodic_long_range(
+        k_grid_mode="fixed", k_grid_bounds=[5, 5, 5], decouple_k_grid=True
+    )
+    before = _module_topologies(long_range, coord, n_node, box, atype)["k_grid"]
+    # a smaller alpha moves the cutoff in; a coupled mask would tighten
+    long_range.kernel.log_alpha = np.log(0.05)
+    after = _module_topologies(long_range, coord, n_node, box, atype)["k_grid"]
+    np.testing.assert_array_equal(after.mask, before.mask)
+    np.testing.assert_array_equal(after.vec, before.vec)
+
+
+def test_decoupled_k_grid_round_trip() -> None:
+    """The frozen cutoff survives serialization, past the kernel drift."""
+    long_range = _periodic_long_range(decouple_k_grid=True)
+    frozen = long_range.frozen_k_max
+    long_range.kernel.log_alpha = np.log(1.5)
+    record = long_range.serialize()
+    assert record["decouple_k_grid"] is True
+    assert record["frozen_k_max"] == pytest.approx(frozen)
+    restored = LatentChargeLongRange.deserialize(record)
+    assert restored.decouple_k_grid
+    assert restored.frozen_k_max == pytest.approx(frozen)
+    # ...even though the restored (drifted) kernel would derive a different one
+    assert restored.kernel.kspace_cutoff(1e-8) != pytest.approx(frozen)
+    # the default stays coupled and serializes nothing
+    plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize())
+    assert not plain.decouple_k_grid
+    assert plain.frozen_k_max is None
+
+
+# ---------------------------------------------------------------------------
+# PME mesh mode
+# ---------------------------------------------------------------------------
+
+
+def _pme_long_range(**kwargs: object) -> LatentChargeLongRange:
+    kwargs.setdefault("k_grid_mode", "pme")
+    return _periodic_long_range(**kwargs)
+
+
+def test_pme_mesh_grid_build_and_flatten() -> None:
+    """The mesh grid packs the cells and the mesh, flags the frames and flattens."""
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        PME_MESH,
+        build_topologies,
+        flatten_topologies,
+        mesh_dims_of_cells,
+        unflatten_topologies,
+    )
+
+    coord = np.random.default_rng(1).normal(size=(4, 3))
+    atype = np.array([0, 1, 1, 0])
+    box = _cells()
+    params = {PME_MESH: {"spacing": 1.0}}
+    tops = build_topologies(
+        ("all_pairs", PME_MESH),
+        coord,
+        np.array([2, 2]),
+        box,
+        atype=atype,
+        topology_params=params,
+    )
+    mesh = tops[PME_MESH]
+    np.testing.assert_array_equal(mesh.box, box)
+    assert mesh.mask.shape == (2,) and mesh.mask.all()
+    dims = tuple(m.shape[0] for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3))
+    assert dims == mesh_dims_of_cells(box, 1.0)
+    # the modes of an axis are its integer frequencies in FFT order
+    n1 = dims[0]
+    expected = np.arange(n1, dtype=float)
+    np.testing.assert_array_equal(
+        mesh.mode_1, np.where(expected > n1 // 2, expected - n1, expected)
+    )
+    assert tops["all_pairs"].index.shape[1] == 4  # guards only
+    flat = flatten_topologies(tops)
+    assert len(flat) == 3 + 8
+    back = unflatten_topologies(("all_pairs", PME_MESH), flat)
+    np.testing.assert_array_equal(back[PME_MESH].box, box)
+    np.testing.assert_array_equal(back[PME_MESH].mask, mesh.mask)
+    np.testing.assert_array_equal(back[PME_MESH].mode_3, mesh.mode_3)
+    # a non-periodic batch: real pairs, masked-out mesh
+    tops_np = build_topologies(
+        ("all_pairs", PME_MESH),
+        coord,
+        np.array([2, 2]),
+        None,
+        atype=atype,
+        topology_params=params,
+    )
+    assert tops_np["all_pairs"].mask.any()
+    assert not tops_np[PME_MESH].mask.any()
+    # missing build parameters are an error
+    with pytest.raises(ValueError, match="spacing"):
+        build_topologies((PME_MESH,), coord, np.array([4]), box, atype=atype)
+
+
+def test_pme_mesh_dims_resolve_every_cell() -> None:
+    """The mesh is the smallest even 7-smooth one resolving every cell."""
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        mesh_dims_of_cells,
+    )
+
+    for spacing in (0.37, 0.5, 1.1):
+        dims = mesh_dims_of_cells(_cells(), spacing)
+        longest = np.max(np.linalg.norm(_cells(), axis=-1), axis=0)
+        for n, a in zip(dims, longest, strict=True):
+            assert n % 2 == 0 and n >= a / spacing
+            m = n
+            for p in (2, 3, 5, 7):
+                while m % p == 0:
+                    m //= p
+            assert m == 1
+            # no smaller even 7-smooth size reaches the spacing
+            for smaller in range(n - 2, 1, -2):
+                m = smaller
+                for p in (2, 3, 5, 7):
+                    while m % p == 0:
+                        m //= p
+                if m == 1:
+                    assert smaller < a / spacing
+                    break
+    # the larger cell alone sizes the mesh of the batch
+    assert mesh_dims_of_cells(_cells(), 0.5) == mesh_dims_of_cells(_cells()[1:], 0.5)
+
+
+def _pme_energy(module, charge, coord, box) -> float:
+    from deepmd.dpmodel.long_range.reciprocal import (
+        mesh_energy,
+    )
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        build_mesh_grid,
+    )
+
+    mesh = build_mesh_grid(box[None], module.mesh_spacing(), coord)
+    frame_id = np.zeros(coord.shape[0], dtype=np.int64)
+    energy = mesh_energy(charge, mesh, frame_id, module.kernel, module.pme_order)
+    return float(np.sum(energy))
+
+
+def test_pme_energy_converges_to_direct_reference() -> None:
+    """The PME energy approaches the direct k-space sum as the mesh refines."""
+    coord = 3.0 + np.random.default_rng(7).uniform(0.0, 1.4, (6, 3))
+    box = np.array([[9.0, 0.4, 0.0], [0.0, 8.5, 0.3], [0.0, 0.0, 9.5]])
+    charge = np.array([[1.0], [-0.5], [1.0], [-0.5], [-0.5], [1.0]])
+    charge = charge - charge.mean(axis=0, keepdims=True)
+
+    kernel = LESKernel(alpha=0.3)
+    # converged direct k-space reference (tol 1e-10)
+    k_max = kernel.kspace_cutoff(1e-10)
+    rec = 2.0 * np.pi * np.linalg.inv(box)
+    n_d = np.ceil(k_max * np.linalg.norm(box, axis=1) / (2 * np.pi)).astype(int)
+    grids = np.meshgrid(*[np.arange(-n, n + 1) for n in n_d], indexing="ij")
+    n_all = np.stack(grids, -1).reshape(-1, 3)
+    k = n_all @ rec.T
+    k2 = (k * k).sum(-1)
+    keep = (k2 > 0) & (k2 <= k_max * k_max)
+    k, k2 = k[keep], k2[keep]
+    phase = k @ coord.T
+    rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge
+    s_sq = (np.abs(rho) ** 2).sum(-1)
+    k_hat = kernel.kspace(k2)
+    volume = abs(np.linalg.det(box))
+    k0 = 2.0 * 0.3 / np.sqrt(np.pi)
+    e_ref = COULOMB_CONSTANT * (
+        (k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum()
+    )
+    errors = []
+    for spacing in (2.0, None):
+        module = LatentChargeLongRange(
+            LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1),
+            kernel,
+            None,
+            periodic=True,
+            k_grid_mode="pme",
+            pme_spacing=spacing,
+        )
+        errors.append(abs(_pme_energy(module, charge, coord, box) - e_ref))
+    assert errors[1] < errors[0]  # refinement improves the energy
+    # the default mesh reaches the accuracy it is chosen for
+    assert errors[1] < 1e-6 * abs(e_ref)
+
+
+def test_pme_mesh_follows_the_cells_of_the_batch() -> None:
+    """The mesh of a batch resolves its cells at the spacing, without statistics."""
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        mesh_dims_of_cells,
+    )
+
+    long_range = _pme_long_range()
+    coord = np.zeros((9, 3))
+    atype = np.zeros(9, dtype=np.int64)
+    small = _module_topologies(long_range, coord, np.array([4, 5]), _cells()[:1], atype)
+    large = _module_topologies(
+        long_range, coord, np.array([4, 5]), 2.0 * _cells(), atype
+    )
+    spacing = long_range.mesh_spacing()
+    assert spacing == pytest.approx(
+        0.75 * np.pi / long_range.kernel.kspace_cutoff(long_range.kspace_tolerance)
+    )
+    for topologies, box in ((small, _cells()[:1]), (large, 2.0 * _cells())):
+        mesh = topologies["pme_mesh"]
+        dims = tuple(m.shape[0] for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3))
+        assert dims == mesh_dims_of_cells(box, spacing)
+        assert mesh.mask.all()
+
+
+def test_pme_mode_configuration() -> None:
+    """The pme mode requires periodicity, and so do the mesh parameters."""
+    with pytest.raises(ValueError, match="periodic"):
+        LatentChargeLongRange(
+            LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1),
+            LESKernel(),
+            None,
+            k_grid_mode="pme",
+        )
+    for kwargs in ({"pme_spacing": 1.0}, {"pme_order": 6}):
+        with pytest.raises(ValueError, match="periodic"):
+            LatentChargeLongRange(
+                LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1),
+                LESKernel(),
+                None,
+                **kwargs,
+            )
+    # every periodic mode takes them: a frozen model evaluates on the mesh
+    assert _periodic_long_range(pme_spacing=1.0).mesh_spacing() == 1.0
+    assert _periodic_long_range().pme_order == 8
+    with pytest.raises(ValueError, match="pme_order"):
+        _pme_long_range(pme_order=1)
+    with pytest.raises(ValueError, match="pme_spacing"):
+        _pme_long_range(pme_spacing=-1.0)
+    with pytest.raises(ValueError, match="k_grid_mode"):
+        _periodic_long_range(k_grid_mode="dense")
+
+
+def test_pme_round_trip() -> None:
+    """The record keeps the pme mode, the order and the spacing."""
+    long_range = _pme_long_range(pme_spacing=1.0, pme_order=5)
+    record = long_range.serialize()
+    assert record["k_grid_mode"] == "pme"
+    assert record["pme_order"] == 5
+    assert record["pme_spacing"] == 1.0
+    restored = LatentChargeLongRange.deserialize(record)
+    assert restored.k_grid_mode == "pme"
+    assert restored.pme_order == 5
+    assert restored.pme_spacing == 1.0
+    # the default order is recorded too, so that it survives a change of default
+    assert _periodic_long_range().serialize()["pme_order"] == 8
diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py
index d147ab3dd3..ad349e95d5 100644
--- a/source/tests/common/dpmodel/test_neighbor_graph.py
+++ b/source/tests/common/dpmodel/test_neighbor_graph.py
@@ -3,6 +3,9 @@
 
 import numpy as np
 
+from deepmd.dpmodel.common import (
+    to_numpy_array,
+)
 from deepmd.dpmodel.utils.neighbor_graph import (
     GraphLayout,
     NeighborGraph,
@@ -11,6 +14,25 @@
 )
 
 
+def assert_csr_consistent(graph: NeighborGraph) -> None:
+    """Check both CSR views against masked endpoints, independently of edge order."""
+    n_nodes = int(to_numpy_array(graph.n_node).sum())
+    mask = to_numpy_array(graph.edge_mask)
+    indices = to_numpy_array(graph.edge_index)
+    for axis, order, row_ptr in (
+        (0, graph.source_order, graph.source_row_ptr),
+        (1, graph.destination_order, graph.destination_row_ptr),
+    ):
+        keys = np.where(mask, indices[axis], n_nodes)
+        np.testing.assert_array_equal(
+            to_numpy_array(order), np.argsort(keys, kind="stable")
+        )
+        counts = np.bincount(indices[axis, mask], minlength=n_nodes)
+        np.testing.assert_array_equal(
+            to_numpy_array(row_ptr), np.concatenate(([0], np.cumsum(counts)))
+        )
+
+
 class TestNeighborGraphDataclass(unittest.TestCase):
     def test_construct_minimal(self) -> None:
         ng = NeighborGraph(
diff --git a/source/tests/common/dpmodel/test_neighbor_preparation_compat.py b/source/tests/common/dpmodel/test_neighbor_preparation_compat.py
new file mode 100644
index 0000000000..383fdb4e3c
--- /dev/null
+++ b/source/tests/common/dpmodel/test_neighbor_preparation_compat.py
@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Keep graph orchestration separate from descriptor preparation semantics."""
+
+import copy
+import unittest
+from unittest.mock import (
+    call,
+    patch,
+)
+
+from deepmd.dpmodel.descriptor.base_descriptor import (
+    BaseDescriptor,
+)
+from deepmd.dpmodel.model.base_model import (
+    BaseModel,
+)
+from deepmd.utils.argcheck import (
+    normalize,
+)
+
+
+def normalized_model(descriptor: dict) -> dict:
+    """Exercise the same strict schema normalization as training input."""
+    return normalize(
+        {
+            "model": {
+                "type_map": ["H"],
+                "descriptor": descriptor,
+                "fitting_net": {},
+            },
+            "learning_rate": {"start_lr": 0.001, "decay_steps": 10},
+            "loss": {},
+            "training": {
+                "numb_steps": 1,
+                "training_data": {"systems": ["unused"]},
+            },
+        },
+        check=True,
+    )["model"]
+
+
+class TestNeighborPreparationCompatibility(unittest.TestCase):
+    """Mock data measurement only, retaining model and descriptor dispatch."""
+
+    @patch("deepmd.utils.update_sel.BaseUpdateSel.get_nbor_stat")
+    def test_omitted_dpa1_sel_keeps_auto_statistics(self, stats) -> None:
+        stats.return_value = (0.5, [25])
+        model = normalized_model({"type": "dpa1", "rcut": 6.0})
+        self.assertEqual(model["descriptor"]["sel"], "auto")
+        before = copy.deepcopy(model)
+        updated, minimum = BaseModel.update_sel(None, ["H"], model)
+        stats.assert_called_once_with(None, ["H"], 6.0, mixed_type=True)
+        self.assertEqual(updated["descriptor"]["sel"], 28)
+        self.assertEqual(minimum, 0.5)
+        self.assertEqual(model, before)
+
+    @patch("deepmd.utils.update_sel.BaseUpdateSel.get_nbor_stat")
+    def test_mixed_hybrid_prepares_and_constructs(self, stats) -> None:
+        stats.side_effect = [(0.5, [4]), (0.4, [4])]
+        model = normalized_model(
+            {
+                "type": "hybrid",
+                "list": [
+                    {
+                        "type": "dpa1",
+                        "rcut": 6.0,
+                        "rcut_smth": 0.5,
+                        "sel": 8,
+                        "neuron": [8, 8],
+                        "axis_neuron": 4,
+                        "attn": 8,
+                        "attn_layer": 0,
+                        "tebd_dim": 4,
+                    },
+                    {
+                        "type": "se_e2_a",
+                        "rcut": 6.0,
+                        "rcut_smth": 0.5,
+                        "sel": [8],
+                        "neuron": [8, 8],
+                        "axis_neuron": 4,
+                    },
+                ],
+            }
+        )
+        updated, minimum = BaseModel.update_sel(None, ["H"], model)
+        self.assertEqual(
+            stats.call_args_list,
+            [
+                call(None, ["H"], 6.0, mixed_type=True),
+                call(None, ["H"], 6.0, mixed_type=False),
+            ],
+        )
+        self.assertEqual(minimum, 0.4)
+        descriptor = BaseDescriptor(**updated["descriptor"], ntypes=1)
+        self.assertEqual(descriptor.get_sel(), [8])
+
+    @patch("deepmd.utils.update_sel.BaseUpdateSel.get_nbor_stat")
+    def test_dpa4c_only_measures_minimum_distance(self, stats) -> None:
+        stats.return_value = (0.5, [0])
+        model = normalized_model({"type": "dpa4c", "rcut": 6.0})
+        before = copy.deepcopy(model)
+        updated, minimum = BaseModel.update_sel(None, ["H"], model)
+        # This tiny cutoff deliberately does not discover neighbor capacity.
+        stats.assert_called_once_with(None, None, 1e-6, mixed_type=True)
+        self.assertEqual(minimum, 0.5)
+        self.assertEqual(updated, before)
+        self.assertEqual(model, before)
+        self.assertNotIn("sel", updated["descriptor"])
+
+    @patch("deepmd.utils.update_sel.BaseUpdateSel.get_nbor_stat")
+    def test_dpa4_preserves_explicit_and_auto_capacity(self, stats) -> None:
+        stats.return_value = (0.5, [25])
+        for sel, expected in [(32, 32), ("auto", 28), ("auto:1.5", 40)]:
+            with self.subTest(sel=sel):
+                stats.reset_mock()
+                model = normalized_model({"type": "dpa4", "rcut": 6.0, "sel": sel})
+                before = copy.deepcopy(model)
+                updated, minimum = BaseModel.update_sel(None, ["H"], model)
+                stats.assert_called_once_with(None, ["H"], 6.0, mixed_type=True)
+                self.assertEqual(updated["descriptor"]["sel"], expected)
+                self.assertEqual(minimum, 0.5)
+                self.assertEqual(model, before)
diff --git a/source/tests/common/dpmodel/test_network.py b/source/tests/common/dpmodel/test_network.py
index 915dc914fa..00d454a782 100644
--- a/source/tests/common/dpmodel/test_network.py
+++ b/source/tests/common/dpmodel/test_network.py
@@ -39,7 +39,7 @@ def test_serialize_deserize(self) -> None:
             ["tanh", "none"],
             [True, False],
             [None, [4], [3, 2]],
-            ["float32", "float64", "single", "double"],
+            ["float32", "float64"],
         ):
             nl0 = NativeLayer(
                 ni,
@@ -181,8 +181,8 @@ def test_embedding_net(self) -> None:
             inp = np.ones([ni], dtype=get_xp_precision(np, prec))
             np.testing.assert_allclose(en0.call(inp), en1.call(inp))
 
-    def test_is_concrete_class(self) -> None:
-        """Verify EmbeddingNet is a concrete class, not factory-generated."""
+    def test_forward_pass(self) -> None:
+        """Test EmbeddingNet forward pass produces correct shapes."""
         in_dim = 4
         neuron = [8, 16, 32]
         net = EmbeddingNet(
@@ -192,27 +192,11 @@ def test_is_concrete_class(self) -> None:
             resnet_dt=True,
             precision="float64",
         )
-        # Check it's the actual EmbeddingNet class, not a dynamic class
-        self.assertEqual(net.__class__.__name__, "EmbeddingNet")
-        self.assertEqual(net.__class__.__module__, "deepmd.dpmodel.utils.network")
-        # Verify it has the expected attributes
         self.assertEqual(net.in_dim, in_dim)
         self.assertEqual(net.neuron, neuron)
         self.assertEqual(net.activation_function, "tanh")
-        self.assertEqual(net.resnet_dt, True)
+        self.assertTrue(net.resnet_dt)
         self.assertEqual(len(net.layers), len(neuron))
-
-    def test_forward_pass(self) -> None:
-        """Test EmbeddingNet forward pass produces correct shapes."""
-        in_dim = 4
-        neuron = [8, 16, 32]
-        net = EmbeddingNet(
-            in_dim=in_dim,
-            neuron=neuron,
-            activation_function="tanh",
-            resnet_dt=True,
-            precision="float64",
-        )
         rng = np.random.default_rng()
         x = rng.standard_normal((5, in_dim))
         out = net.call(x)
@@ -310,37 +294,8 @@ def test_fitting_net(self) -> None:
             )
             en1 = FittingNet.deserialize(en0.serialize())
             inp = np.ones([ni], dtype=get_xp_precision(np, prec))
-            en0.call(inp)
-            en1.call(inp)
             np.testing.assert_allclose(en0.call(inp), en1.call(inp))
 
-    def test_is_concrete_class(self) -> None:
-        """Verify FittingNet is a concrete class, not factory-generated."""
-        in_dim = 4
-        out_dim = 1
-        neuron = [8, 16]
-        net = FittingNet(
-            in_dim=in_dim,
-            out_dim=out_dim,
-            neuron=neuron,
-            activation_function="tanh",
-            resnet_dt=True,
-            precision="float64",
-            bias_out=True,
-        )
-        # Check it's the actual FittingNet class, not a dynamic class
-        self.assertEqual(net.__class__.__name__, "FittingNet")
-        self.assertEqual(net.__class__.__module__, "deepmd.dpmodel.utils.network")
-        # Verify it has the expected attributes
-        self.assertEqual(net.in_dim, in_dim)
-        self.assertEqual(net.out_dim, out_dim)
-        self.assertEqual(net.neuron, neuron)
-        self.assertEqual(net.activation_function, "tanh")
-        self.assertEqual(net.resnet_dt, True)
-        self.assertEqual(net.bias_out, True)
-        # FittingNet has len(neuron) embedding layers + 1 output layer
-        self.assertEqual(len(net.layers), len(neuron) + 1)
-
     def test_forward_pass(self) -> None:
         """Test FittingNet forward pass produces correct output shape."""
         in_dim = 4
@@ -354,6 +309,13 @@ def test_forward_pass(self) -> None:
             resnet_dt=True,
             precision="float64",
         )
+        self.assertEqual(net.in_dim, in_dim)
+        self.assertEqual(net.neuron, neuron)
+        self.assertEqual(net.activation_function, "tanh")
+        self.assertTrue(net.resnet_dt)
+        self.assertEqual(net.out_dim, out_dim)
+        self.assertTrue(net.bias_out)
+        self.assertEqual(len(net.layers), len(neuron) + 1)
         # Single sample
         rng = np.random.default_rng()
         x = rng.standard_normal(in_dim)
diff --git a/source/tests/common/dpmodel/test_spin_model_virtual_types.py b/source/tests/common/dpmodel/test_spin_model_virtual_types.py
index 727eac3b53..4795dfb971 100644
--- a/source/tests/common/dpmodel/test_spin_model_virtual_types.py
+++ b/source/tests/common/dpmodel/test_spin_model_virtual_types.py
@@ -80,6 +80,18 @@ def test_dense_spin_expansion_preserves_virtual_types(model) -> None:
     )
 
 
+def test_real_type_sequence_bias_override(model) -> None:
+    sample = {
+        "coord": np.arange(9, dtype=np.float64).reshape(1, 3, 3),
+        "atype": np.array([[0, 1, 2]], dtype=np.int32),
+        "spin": np.ones((1, 3, 3), dtype=np.float64),
+    }
+    model.change_out_bias([sample], preset_out_bias={"energy": [1.0, 2.0, 3.0]})
+    np.testing.assert_allclose(
+        to_numpy_array(model.get_out_bias()).reshape(-1), [1.0, 2.0, 3.0, 0.0, 0.0, 0.0]
+    )
+
+
 def test_lower_spin_expansion_preserves_virtual_types(model) -> None:
     """Local and ghost placeholders remain negative in the switched layout."""
     extended_coord = np.arange(12, dtype=np.float64).reshape(1, 4, 3)
diff --git a/source/tests/common/dpmodel/test_train_schedule.py b/source/tests/common/dpmodel/test_train_schedule.py
index dd85991f2c..b4feda6a9e 100644
--- a/source/tests/common/dpmodel/test_train_schedule.py
+++ b/source/tests/common/dpmodel/test_train_schedule.py
@@ -171,6 +171,22 @@ def test_multi_task_steps_keep_configured_model_prob() -> None:
     np.testing.assert_allclose(schedule.model_prob, [0.75, 0.25])
 
 
+@pytest.mark.parametrize("model_prob", [None, {}])
+def test_default_multi_task_probabilities_use_system_counts(
+    model_prob: dict[str, float] | None,
+) -> None:
+    schedule = resolve_step_schedule(
+        {"numb_steps": 12, "model_prob": model_prob},
+        multi_task=True,
+        model_keys=["small", "large"],
+        training_data={"small": FakeData(1), "large": FakeData(3)},
+        epoch_length=_unreachable_epoch_length,
+    )
+
+    assert schedule.num_steps == 12
+    np.testing.assert_allclose(schedule.model_prob, [0.25, 0.75])
+
+
 def test_multi_task_requires_a_run_length() -> None:
     with pytest.raises(ValueError, match="num_epoch_dict must be set"):
         resolve_step_schedule(
diff --git a/source/tests/common/dpmodel/test_update_sel.py b/source/tests/common/dpmodel/test_update_sel.py
index ff8313b204..5f11d511d4 100644
--- a/source/tests/common/dpmodel/test_update_sel.py
+++ b/source/tests/common/dpmodel/test_update_sel.py
@@ -143,33 +143,6 @@ def test_update_sel_atten_int(self, sel_mock) -> None:
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat")
-    def test_update_sel_atten_list(self, sel_mock) -> None:
-        sel_mock.return_value = self.mock_min_nbor_dist, [25]
-
-        jdata = {
-            "model": {
-                "descriptor": {
-                    "type": "se_atten",
-                    "sel": 30,
-                    "rcut": 6,
-                }
-            },
-            "training": {"training_data": {}},
-        }
-        expected_out = {
-            "model": {
-                "descriptor": {
-                    "type": "se_atten",
-                    "sel": 30,
-                    "rcut": 6,
-                }
-            },
-            "training": {"training_data": {}},
-        }
-        jdata = update_sel(jdata)
-        self.assertEqual(jdata, expected_out)
-
     def test_wrap_up_4(self) -> None:
         self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4)
         self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4)
diff --git a/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py b/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
new file mode 100644
index 0000000000..365185f8e7
--- /dev/null
+++ b/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
@@ -0,0 +1,392 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Graph-native isolated references must not allocate a dense sel-sized list."""
+
+import dataclasses
+from unittest import (
+    mock,
+)
+
+import numpy as np
+import pytest
+
+from deepmd.dpmodel.atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.dpmodel.descriptor.dpa1 import (
+    DescrptDPA1,
+)
+from deepmd.dpmodel.descriptor.dpa2 import (
+    DescrptDPA2,
+)
+from deepmd.dpmodel.descriptor.dpa4 import (
+    DescrptDPA4,
+)
+from deepmd.dpmodel.descriptor.dpa4c import (
+    DescrptDPA4C,
+)
+from deepmd.dpmodel.fitting import (
+    InvarFitting,
+)
+from deepmd.dpmodel.utils.neighbor_graph import (
+    graph_from_dense_quartet,
+)
+
+from ...dpa4_fixtures import (
+    jitter_zero_arrays,
+)
+
+TYPE_MAP = ["O", "H"]
+BIAS = [[-3.0], [0.5]]
+
+
+def make_model(
+    kind: str,
+    charge: bool = False,
+    spin: bool = False,
+    precision: str = "float64",
+    numb_fparam: int = 0,
+    numb_aparam: int = 0,
+) -> DPAtomicModel:
+    """Activate learned condition heads so missing inputs cannot pass by chance."""
+    if kind == "dpa1":
+        descriptor = DescrptDPA1(
+            rcut=4.0,
+            rcut_smth=0.5,
+            sel=[2],
+            ntypes=2,
+            attn_layer=0,
+            axis_neuron=2,
+            neuron=[4, 8],
+            precision=precision,
+            seed=3,
+        )
+    elif kind == "dpa2":
+        descriptor = DescrptDPA2(
+            ntypes=2,
+            repinit={
+                "rcut": 4.0,
+                "rcut_smth": 0.5,
+                "nsel": 2,
+                "neuron": [4, 8],
+                "axis_neuron": 2,
+            },
+            repformer={
+                "rcut": 3.0,
+                "rcut_smth": 0.5,
+                "nsel": 1,
+                "nlayers": 1,
+                "g1_dim": 8,
+                "g2_dim": 4,
+                "axis_neuron": 2,
+                "update_g1_has_attn": False,
+            },
+            precision=precision,
+            seed=3,
+        )
+    else:
+        cls = DescrptDPA4 if kind == "dpa4" else DescrptDPA4C
+        kwargs = {
+            "ntypes": 2,
+            "rcut": 4.0,
+            "channels": 8,
+            "n_radial": 4,
+            "lmax": 2,
+            "precision": precision,
+            "seed": 3,
+            "add_chg_spin_ebd": charge,
+            "default_chg_spin": [0, 1] if charge else None,
+            "use_spin": [True, False] if spin else None,
+        }
+        if kind == "dpa4":
+            kwargs.update(sel=8, mmax=1, n_blocks=1)
+        descriptor = cls(**kwargs)
+        descriptor = cls.deserialize(
+            jitter_zero_arrays(descriptor.serialize(), np.random.default_rng(7))
+        )
+        if kind == "dpa4c" and spin:
+            descriptor.spin.spin_gate[...] = 1.0
+    fitting = InvarFitting(
+        "energy",
+        2,
+        descriptor.get_dim_out(),
+        1,
+        neuron=[8],
+        mixed_types=descriptor.mixed_types(),
+        precision=precision,
+        vacuum_ref=True,
+        numb_fparam=numb_fparam,
+        numb_aparam=numb_aparam,
+        seed=1,
+    )
+    return DPAtomicModel(
+        descriptor,
+        fitting,
+        type_map=TYPE_MAP,
+        preset_out_bias={"energy": BIAS},
+    )
+
+
+def bounded_reference(model: DPAtomicModel):
+    """Use two explicitly bounded dense slots, never the descriptor's sel.
+
+    The native-spin interface is graph-only, so adapt the dense quartet and
+    pass spin on its flat node axis. Spin-less models also exercise the
+    descriptor's dense compatibility adapter directly.
+    """
+    dtype = model.fitting_net.bias_atom_e.dtype
+    coord = np.zeros((2, 1, 3), dtype=dtype)
+    atype = np.arange(2, dtype=np.int64).reshape(2, 1)
+    nlist = np.full((2, 1, 2), -1, dtype=np.int64)
+    mapping = np.zeros((2, 1), dtype=np.int64)
+    graph, flat_atype = graph_from_dense_quartet(coord, atype, nlist, mapping)
+    conditions = {
+        key: np.asarray(value, dtype=dtype)
+        for key, value in model.vacuum_conditions().items()
+    }
+    descriptor = model.descriptor
+    expected = descriptor.call_graph(
+        graph,
+        flat_atype,
+        type_embedding=descriptor.graph_type_embedding_table(),
+        **conditions,
+    )[0]
+    if not model.add_spin_ebd:
+        dense = descriptor(coord, atype, nlist, mapping=mapping, **conditions)[0]
+        rtol, atol = (1e-7, 2e-6) if dtype == np.float32 else (1e-12, 1e-12)
+        np.testing.assert_allclose(
+            dense.reshape(expected.shape), expected, rtol=rtol, atol=atol
+        )
+    return expected, graph, flat_atype, conditions
+
+
+@pytest.mark.parametrize("kind", ["dpa4", "dpa4c"])
+@pytest.mark.parametrize(
+    "charge,spin", [(False, False), (True, False), (False, True), (True, True)]
+)
+@pytest.mark.parametrize("precision", ["float32", "float64"])
+def test_vacuum_uses_bounded_graph(kind, charge, spin, precision) -> None:
+    tol = 2e-6 if precision == "float32" else 1e-12
+    model = make_model(kind, charge, spin, precision)
+    expected, _, _, conditions = bounded_reference(model)
+    descriptor = model.descriptor
+    assert descriptor.supports_graph_vacuum()
+    if kind == "dpa4c":
+        assert descriptor.get_sel() == [999999]
+    call_graph = descriptor.call_graph
+
+    def checked_graph(graph, atype, **kwargs):
+        np.testing.assert_array_equal(graph.n_node, [1, 1])
+        np.testing.assert_array_equal(atype, [0, 1])
+        assert graph.edge_index.shape == (2, 2)
+        assert graph.edge_vec.shape == (2, 3)
+        assert graph.edge_mask.shape == (2,)
+        assert not graph.edge_mask.any()
+        assert graph.edge_vec.dtype == model.fitting_net.bias_atom_e.dtype
+        assert set(kwargs) == {"type_embedding", *conditions}
+        for name, value in conditions.items():
+            np.testing.assert_array_equal(kwargs[name], value)
+        return call_graph(graph, atype, **kwargs)
+
+    # Fail before allocation if the implementation attempts the dense route.
+    # Repeat with a larger sentinel to prove that the graph is sel-independent.
+    with (
+        mock.patch.object(
+            descriptor, "get_nnei", side_effect=AssertionError("dense capacity queried")
+        ),
+        mock.patch.object(
+            descriptor, "call", side_effect=AssertionError("dense adapter called")
+        ),
+        mock.patch.object(
+            descriptor, "call_graph", side_effect=checked_graph
+        ) as graph_call,
+    ):
+        for sentinel in (999999, 1999999):
+            with mock.patch.object(descriptor, "get_sel", return_value=[sentinel]):
+                actual = model.vacuum_descriptor()
+            assert actual.shape == (2, descriptor.get_dim_out())
+            assert actual.dtype == model.fitting_net.bias_atom_e.dtype
+            np.testing.assert_allclose(actual, expected, rtol=tol, atol=tol)
+        assert graph_call.call_count == 2
+    assert np.max(np.abs(actual)) > 1e-5
+    assert not np.allclose(actual[0], actual[1])
+
+
+@pytest.mark.parametrize("kind", ["dpa4", "dpa4c"])
+def test_empty_and_masked_edges_preserve_conditions(kind) -> None:
+    model = make_model(kind, charge=True, spin=True)
+    expected, dense_graph, atype, conditions = bounded_reference(model)
+    np.testing.assert_allclose(model.vacuum_descriptor(), expected, atol=1e-12)
+    for nedges in (0, 1, 7):
+        graph = dataclasses.replace(
+            dense_graph,
+            edge_index=np.tile(np.array([[1], [0]], dtype=np.int64), (1, nedges)),
+            edge_vec=np.ones((nedges, 3)),
+            edge_mask=np.zeros(nedges, dtype=bool),
+            destination_order=None,
+            destination_row_ptr=None,
+            source_order=None,
+            source_row_ptr=None,
+        )
+        actual = model.descriptor.call_graph(graph, atype, **conditions)[0]
+        np.testing.assert_allclose(actual, expected, atol=1e-12)
+    # The reference is the learned, ground-state-conditioned value, not zero
+    # or a default-condition shortcut, even when no real edges exist.
+    for name in conditions:
+        other = {**conditions, name: np.zeros_like(conditions[name])}
+        if name == "charge_spin":
+            other[name][:, 1] = 1.0
+        actual = model.descriptor.call_graph(dense_graph, atype, **other)[0]
+        assert np.max(np.abs(actual - expected)) > 1e-7, name
+
+
+def test_graph_descriptor_receives_its_type_embedding() -> None:
+    """An opted-in descriptor receives its external type-embedding table."""
+    model = make_model("dpa1")
+    assert not model.descriptor.supports_graph_vacuum()
+    expected, _, _, _ = bounded_reference(model)
+    with (
+        mock.patch.object(model.descriptor, "supports_graph_vacuum", return_value=True),
+        mock.patch.object(
+            model.descriptor, "call", side_effect=AssertionError("dense adapter called")
+        ),
+    ):
+        actual = model.vacuum_descriptor()
+    np.testing.assert_allclose(actual, expected, atol=1e-12)
+
+
+def test_dense_forward_keeps_its_own_reference_semantics() -> None:
+    """DPA1's legacy padding residual must cancel on its dense path too."""
+    model = make_model("dpa1")
+    model.descriptor.se_atten.mean[...] = 0.3
+    coord = np.zeros((2, 1, 3))
+    atype = np.arange(2, dtype=np.int64).reshape(2, 1)
+    nlist = np.full((2, 1, 2), -1, dtype=np.int64)
+    mapping = np.zeros((2, 1), dtype=np.int64)
+    assert model.descriptor.uses_graph_lower()
+    assert not model.descriptor.supports_graph_vacuum()
+    graph, flat_atype = graph_from_dense_quartet(coord, atype, nlist, mapping)
+    graph_reference = model.descriptor.call_graph(
+        graph,
+        flat_atype,
+        type_embedding=model.descriptor.graph_type_embedding_table(),
+    )[0]
+    dense_reference = model.vacuum_descriptor()
+    assert np.max(np.abs(graph_reference - dense_reference)) > 1e-7
+    before = model.forward_atomic(coord, atype, nlist, mapping=mapping)["energy"]
+    np.testing.assert_allclose(
+        before[:, 0], model.fitting_net.bias_atom_e, rtol=1e-12, atol=1e-12
+    )
+    # Graph support alone must not silently change the legacy reference.
+    model.fold_vacuum_reference()
+    after = model.forward_atomic(coord, atype, nlist, mapping=mapping)["energy"]
+    np.testing.assert_allclose(after, before, rtol=1e-12, atol=1e-12)
+
+
+def test_disabled_graph_lower_keeps_dense_reference() -> None:
+    """An explicit lower opt-out still selects the dense compatibility path."""
+    model = make_model("dpa4", charge=True)
+    expected = model.vacuum_descriptor()
+    model.descriptor.disable_graph_lower()
+    with mock.patch.object(
+        model.descriptor, "call", wraps=model.descriptor.call
+    ) as call:
+        actual = model.vacuum_descriptor()
+    call.assert_called_once()
+    np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize(
+    "kind,charge,spin",
+    [
+        ("dpa4", False, False),
+        ("dpa4", True, True),
+        ("dpa4c", False, False),
+        ("dpa4c", True, True),
+    ],
+)
+def test_folding_matches_appended_reference_frames(kind, charge, spin) -> None:
+    """Folding must match reference nodes carried alongside a real cluster."""
+    model = make_model(kind, charge, spin)
+    coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.2, 0.0]]])
+    atype = np.array([[0, 1]], dtype=np.int64)
+    nlist = np.array([[[1, -1], [0, -1]]], dtype=np.int64)
+    graph, atype = graph_from_dense_quartet(coord, atype, nlist, None)
+    conditions = {}
+    if charge:
+        conditions["charge_spin"] = np.array([[1.0, 2.0]])
+    if spin:
+        conditions["spin"] = np.array([[0.2, 0.4, 1.0], [0.0, 0.0, 0.0]])
+    before = model.forward_atomic_graph(graph, atype, **conditions)["energy"]
+    model.fold_vacuum_reference()
+    assert not model.fitting_net.needs_vacuum_descriptor()
+    after = model.forward_atomic_graph(graph, atype, **conditions)["energy"]
+    np.testing.assert_allclose(after, before, rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize("kind", ["dpa1", "dpa2"])  # legacy descriptor
+@pytest.mark.parametrize(
+    "numb_fparam,numb_aparam", [(0, 0), (1, 0), (0, 1)]
+)  # bias folding or conditioned reference table
+def test_legacy_graph_uses_dense_vacuum_reference(
+    kind, numb_fparam, numb_aparam
+) -> None:
+    """Graph inference and folding must subtract the same legacy reference."""
+    model = make_model(kind, numb_fparam=numb_fparam, numb_aparam=numb_aparam)
+    descriptor = model.descriptor
+    blocks = (
+        [descriptor.se_atten]
+        if kind == "dpa1"
+        else [descriptor.repinit, descriptor.repformers]
+    )
+    for block in blocks:
+        block.mean[...] = 0.3
+    assert descriptor.uses_graph_lower()
+    assert not descriptor.supports_graph_vacuum()
+    iso_coord = np.zeros((2, 1, 3))
+    iso_atype = np.arange(2, dtype=np.int64).reshape(2, 1)
+    iso_nlist = np.full((2, 1, 2), -1, dtype=np.int64)
+    iso_graph, iso_atype = graph_from_dense_quartet(
+        iso_coord, iso_atype, iso_nlist, None
+    )
+    graph_reference = descriptor.call_graph(
+        iso_graph, iso_atype, type_embedding=descriptor.graph_type_embedding_table()
+    )[0]
+    dense_reference = model.vacuum_descriptor()
+    assert np.max(np.abs(graph_reference - dense_reference)) > 1e-7
+
+    coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.2, 0.0]]])
+    atype = np.array([[0, 1]], dtype=np.int64)
+    nlist = np.array([[[1, -1], [0, -1]]], dtype=np.int64)
+    graph, atype = graph_from_dense_quartet(coord, atype, nlist, None)
+    fparam = np.array([[0.7]]) if numb_fparam else None
+    aparam = np.array([[0.2], [0.6]]) if numb_aparam else None
+    gg = descriptor.call_graph(
+        graph, atype, type_embedding=descriptor.graph_type_embedding_table()
+    )[0]
+    expected = model.fitting_net.call_graph(
+        gg,
+        atype,
+        fparam=np.repeat(fparam, 2, axis=0) if numb_fparam else None,
+        aparam=aparam,
+        vacuum_descriptor=dense_reference,
+    )["energy"]
+    with (
+        mock.patch.object(
+            model, "append_vacuum_frames", wraps=model.append_vacuum_frames
+        ) as append,
+        mock.patch.object(
+            model, "vacuum_descriptor", wraps=model.vacuum_descriptor
+        ) as reference,
+    ):
+        before = model.forward_atomic_graph(graph, atype, fparam=fparam, aparam=aparam)[
+            "energy"
+        ]
+    np.testing.assert_allclose(before, expected, rtol=1e-12, atol=1e-12)
+    append.assert_not_called()
+    reference.assert_called_once()
+    model.fold_vacuum_reference()
+    assert not model.fitting_net.needs_vacuum_descriptor()
+    after = model.forward_atomic_graph(graph, atype, fparam=fparam, aparam=aparam)[
+        "energy"
+    ]
+    np.testing.assert_allclose(after, before, rtol=1e-12, atol=1e-12)
diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py
index 9b61ee26cf..061ada901a 100644
--- a/source/tests/common/dpmodel/test_zbl_bridging.py
+++ b/source/tests/common/dpmodel/test_zbl_bridging.py
@@ -606,6 +606,9 @@ def has_default_chg_spin(self) -> bool:
         def get_default_chg_spin(self):
             return None
 
+        def has_long_range(self) -> bool:
+            return False
+
     def _compose(self, children):
         return LinearEnergyAtomicModel(children, type_map=["Ni", "O"], weights="sum")
 
diff --git a/source/tests/common/sezm_config.py b/source/tests/common/sezm_config.py
new file mode 100644
index 0000000000..a4d50d1965
--- /dev/null
+++ b/source/tests/common/sezm_config.py
@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Backend-independent SeZM configurations for export tests."""
+
+
+def _tiny_sezm_model_params() -> dict:
+    """Minimal fp64 SeZM config for self-contained export tests.
+
+    ``precision="float64"`` is what unlocks the ``rtol=1e-10, atol=1e-10``
+    parity pt_expt enforces; fp32 accumulation alone drifts in the 1e-6
+    range. All other knobs keep ``make_fx`` tracing small.
+    """
+    return {
+        "type": "SeZM",
+        "type_map": ["A", "B"],
+        "descriptor": {
+            "type": "SeZM",
+            "sel": [2, 2],
+            "rcut": 3.0,
+            "channels": 4,
+            "n_focus": 1,
+            "n_radial": 3,
+            "radial_mlp": [6],
+            "use_env_seed": True,
+            "l_schedule": [1, 0],
+            "mmax": 1,
+            "so2_norm": False,
+            "so2_layers": 1,
+            "n_atten_head": 1,
+            "sandwich_norm": [True, False, True, False],
+            "ffn_neurons": 8,
+            "ffn_blocks": 1,
+            "s2_activation": [False, True],
+            "mlp_bias": False,
+            "layer_scale": False,
+            "use_amp": False,
+            "activation_function": "silu",
+            "glu_activation": True,
+            "precision": "float64",
+            "seed": 7,
+        },
+        "fitting_net": {
+            "neuron": [8],
+            "activation_function": "silu",
+            "precision": "float64",
+            "seed": 7,
+        },
+        "use_compile": False,
+    }
+
+
+def _tiny_sezm_native_spin_model_params() -> dict:
+    """Minimal fp64 native-spin SeZM config for freeze routing tests."""
+    params = _tiny_sezm_model_params()
+    params["type_map"] = ["O", "H"]
+    params["spin"] = {
+        "use_spin": [True, False],
+        "scheme": "native",
+    }
+    return params
diff --git a/source/tests/common/test_agent_skills.py b/source/tests/common/test_agent_skills.py
index daab80ab6e..b5e4131873 100644
--- a/source/tests/common/test_agent_skills.py
+++ b/source/tests/common/test_agent_skills.py
@@ -12,21 +12,7 @@
 
 ROOT = Path(__file__).resolve().parents[3]
 LAMMPS_SKILL = ROOT / "skills" / "lammps-deepmd" / "SKILL.md"
-LAMMPS_DEPLOYMENT = (
-    ROOT / "skills" / "lammps-deepmd" / "references" / "model-deployment.md"
-)
-LAMMPS_WORKFLOW = (
-    ROOT / "skills" / "lammps-deepmd" / "references" / "commands-and-workflow.md"
-)
 LAMMPS_ASSET = ROOT / "skills" / "lammps-deepmd" / "assets" / "input.nvt.lammps"
-DPA4_TRAIN_REFERENCE = ROOT / "skills" / "deepmd-train" / "models" / "dpa4.md"
-DPA4C_TRAIN_REFERENCE = ROOT / "skills" / "deepmd-train" / "models" / "dpa4c.md"
-DPA4C_MODEL_DOC = ROOT / "doc" / "model" / "dpa4c.md"
-DPA4_FREEZE_POLICY = (
-    ROOT / "skills" / "deepmd-python-inference" / "references" / "dpa4-freeze-policy.md"
-)
-DPA4_FINETUNE_SKILL = ROOT / "skills" / "deepmd-finetune-dpa4" / "SKILL.md"
-INFERENCE_SKILL = ROOT / "skills" / "deepmd-python-inference" / "SKILL.md"
 HELD_OUT_REFERENCE = (
     ROOT
     / "skills"
@@ -34,89 +20,10 @@
     / "references"
     / "held-out-evaluation.md"
 )
-DP_TEST_ENTRYPOINT = ROOT / "deepmd" / "entrypoints" / "test.py"
-ENERGY_TESTER = ROOT / "deepmd" / "infer" / "model_test" / "ener.py"
-FINETUNE_SOURCE = ROOT / "deepmd" / "utils" / "finetune.py"
-
-
-def test_lammps_skill_uses_capability_gated_runtime() -> None:
-    text = LAMMPS_SKILL.read_text(encoding="utf-8")
-
-    assert "3.2.0b0" not in text
-    assert "uvx --from" not in text
-    assert "record the resolved Git commit SHA" in text
-    assert "do not claim support from an" in text
-    assert "unreleased version number" in text
-    assert "do not install or upgrade packages silently" in text
-
-
-def test_lammps_required_inputs_remain_nested() -> None:
-    text = LAMMPS_SKILL.read_text(encoding="utf-8")
-    section = text.split("1. Confirm the minimum simulation inputs:", 1)[1].split(
-        "1. Write the LAMMPS input script", 1
-    )[0]
-
-    for required in (
-        "structure/data file",
-        "DeePMD model artifact",
-        "atom type to element mapping",
-        "target ensemble",
-        "temperature, pressure",
-    ):
-        assert any(
-            line.startswith("   - ") and required in line
-            for line in section.splitlines()
-        )
-
-
-def test_complete_held_out_evaluation_is_routed_and_evidence_complete() -> None:
-    finetune = DPA4_FINETUNE_SKILL.read_text(encoding="utf-8")
-    inference = INFERENCE_SKILL.read_text(encoding="utf-8")
-    held_out = HELD_OUT_REFERENCE.read_text(encoding="utf-8")
-
-    assert "references/held-out-evaluation.md" in inference
-    assert "held-out-evaluation.md" in finetune
-    assert "dp --pt test -m selected.pt" in held_out
-    assert '-n 0 -d "$detail_prefix"' in held_out
-    assert "one command per held-out system" in held_out
-    assert "population standard deviation (`ddof=0`)" in held_out
-    assert "do not average per-system RMSE values" in held_out
-    assert "Training logs, a successful freeze" in held_out
-
-
-def test_complete_held_out_evaluation_checks_dataset_shapes() -> None:
-    held_out = HELD_OUT_REFERENCE.read_text(encoding="utf-8")
-
-    assert "number of whitespace-separated entries in `type.raw`" in held_out
-    assert "coordinate and force widths of `3 * natoms`" in held_out
-    assert "force rows equal `frames * natoms`" in held_out
-    assert "divided by `natoms` exactly once" in held_out
-    assert "zero-based indices" in held_out
-
-
-def test_held_out_contract_matches_dp_test_source() -> None:
-    entrypoint = DP_TEST_ENTRYPOINT.read_text(encoding="utf-8")
-    energy_tester = ENERGY_TESTER.read_text(encoding="utf-8")
-
-    assert "if numb_test == 0:" in entrypoint
-    assert 'detail_path.with_suffix(".e.out")' in energy_tester
-    assert 'detail_path.with_suffix(".e_peratom.out")' in energy_tester
-    assert 'detail_path.with_suffix(".f.out")' in energy_tester
-
-
-def test_dpa4_freeze_and_lammps_stay_on_target_node() -> None:
-    deployment = LAMMPS_DEPLOYMENT.read_text(encoding="utf-8")
-
-    assert "same target physical compute node" in deployment
-    assert "inspect the native checkpoint -> freeze `.pt2` -> `run 0`" in deployment
-    assert "move the archive to B" in deployment
-    assert "artifact portability is proven" in deployment
 
 
 def test_lammps_mapping_data_and_dump_contract() -> None:
     skill = LAMMPS_SKILL.read_text(encoding="utf-8")
-    deployment = LAMMPS_DEPLOYMENT.read_text(encoding="utf-8")
-    workflow = LAMMPS_WORKFLOW.read_text(encoding="utf-8")
     example = skill.split("## Example: annotated NVT input", 1)[1].split(
         "### What every command means", 1
     )[0]
@@ -126,24 +33,6 @@ def test_lammps_mapping_data_and_dump_contract() -> None:
     )
     assert "pair_coeff      * * Si O" in example
     assert "dump_modify     1 element Si O sort id" in example
-    assert "zero-based `type.raw` index" in deployment
-    assert "Do not sort atoms by element" in deployment
-    assert "`xy = b_x`" in deployment
-    assert "`xz = c_x`" in deployment
-    assert "`yz = c_y`" in deployment
-    assert "first line of a LAMMPS data file is a title" in deployment
-    assert "first line of a LAMMPS data file is a skipped title" in workflow
-
-
-def test_lammps_canary_requires_physical_stability() -> None:
-    deployment = LAMMPS_DEPLOYMENT.read_text(encoding="utf-8")
-    workflow = LAMMPS_WORKFLOW.read_text(encoding="utf-8")
-
-    assert "short NVE when physically appropriate" in deployment
-    assert "Exit code zero alone is not a passed canary" in deployment
-    assert "early temperature, pressure, and controlled variables" in deployment
-    assert "do not barostat that direction" in workflow
-    assert "shell environment variables and LAMMPS variables distinct" in workflow
 
 
 def test_held_out_command_runs_from_clean_directory_without_overwrite(
@@ -254,23 +143,6 @@ def test_dpa4_and_dpa4c_preset_examples_expand_and_validate() -> None:
         assert normalized["model"]["fitting_net"]["type"] == expected_fitting_type
 
 
-def test_dpa4c_skill_routes_backend_specific_compile_options() -> None:
-    top_level = (ROOT / "skills" / "deepmd-train" / "SKILL.md").read_text(
-        encoding="utf-8"
-    )
-    reference = DPA4C_TRAIN_REFERENCE.read_text(encoding="utf-8")
-
-    assert "models/dpa4c.md" in top_level
-    assert "dp --pt-expt train input.json" in reference
-    assert "training.enable_compile" in reference
-    assert "training.enable_tf32" in reference
-    assert "model.use_compile" in reference
-    assert "Do not put `use_compile` under `model` for DPA4C" in reference
-    assert "preset_out_bias" in reference
-    assert "vacuum_ref" in reference
-    assert "disables" in reference
-
-
 def test_dpa4c_reference_configuration_initializes_trainer(
     tmp_path: Path,
 ) -> None:
@@ -303,22 +175,6 @@ def test_dpa4c_reference_configuration_initializes_trainer(
     assert trainer.lr_schedule is not None
 
 
-def test_dpa4c_export_policy_matches_model_documentation() -> None:
-    reference = DPA4C_TRAIN_REFERENCE.read_text(encoding="utf-8")
-    model_doc = DPA4C_MODEL_DOC.read_text(encoding="utf-8")
-    deployment = LAMMPS_DEPLOYMENT.read_text(encoding="utf-8")
-
-    for variable in ("DP_CUDA_INFER", "DP_TF32_INFER", "DP_AMP_INFER"):
-        assert f"export {variable}=" in reference
-        assert variable in model_doc
-    assert "`DP_CUDA_INFER` of at least `1`" in reference
-    assert "`DP_CUDA_INFER` of at least `1`" in model_doc
-    assert "same target physical compute node" in reference
-    assert "freeze `.pt2` -> compress `.pt2`" in deployment
-    assert "dp --pt-expt freeze" in deployment
-    assert "dp --pt_expt" not in deployment
-
-
 def test_lammps_asset_matches_mapping_contract() -> None:
     asset = LAMMPS_ASSET.read_text(encoding="utf-8")
 
@@ -328,41 +184,3 @@ def test_lammps_asset_matches_mapping_contract() -> None:
     assert "pair_coeff      * * Si O" in asset
     assert "id type element x y z" in asset
     assert "dump_modify     1 element Si O sort id" in asset
-
-
-def test_dpa4_freeze_policy_is_explicit_and_routed() -> None:
-    policy = DPA4_FREEZE_POLICY.read_text(encoding="utf-8")
-    train = DPA4_TRAIN_REFERENCE.read_text(encoding="utf-8")
-    finetune = DPA4_FINETUNE_SKILL.read_text(encoding="utf-8")
-    deployment = LAMMPS_DEPLOYMENT.read_text(encoding="utf-8")
-
-    for variable in ("DP_TRITON_INFER", "DP_TF32_INFER", "DP_AMP_INFER"):
-        assert f"export {variable}=" in policy
-    assert "Levels 1 and 2 keep" in policy
-    assert "Level 3, TF32, and AMP" in policy
-    assert "dpa4-freeze-policy.md" in train
-    assert "dpa4-freeze-policy.md" in finetune
-    assert "dpa4-freeze-policy.md" in deployment
-
-
-def test_held_out_multitask_and_optional_type_map_contracts() -> None:
-    held_out = HELD_OUT_REFERENCE.read_text(encoding="utf-8")
-
-    assert "When `type_map.raw` is present" in held_out
-    assert "is absent, require provenance" in held_out
-    assert "--head SELECTED_BRANCH" in held_out
-    assert "already single-head; do not pass `--head`" in held_out
-
-
-def test_use_pretrain_script_guidance_matches_source_scope() -> None:
-    guidance = DPA4_FINETUNE_SKILL.read_text(encoding="utf-8")
-    source = FINETUNE_SOURCE.read_text(encoding="utf-8")
-    function = source.split("def _apply_pretrained_model_params", 1)[1].split(
-        "\ndef ", 1
-    )[0]
-
-    assert "does not restore the complete model" in guidance
-    assert "model.descriptor" in guidance
-    assert "model.fitting_net" in guidance
-    assert 'pretrained_config["descriptor"]' in function
-    assert 'pretrained_config["fitting_net"]' in function
diff --git a/source/tests/common/test_argcheck_training.py b/source/tests/common/test_argcheck_training.py
index cb5d2ffafa..b954aedcec 100644
--- a/source/tests/common/test_argcheck_training.py
+++ b/source/tests/common/test_argcheck_training.py
@@ -1,7 +1,19 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Tests for normalized training arguments."""
 
+import copy
+from typing import (
+    Any,
+)
+
+import pytest
+from dargs.dargs import (
+    ArgumentError,
+    ArgumentTypeError,
+)
+
 from deepmd.utils.argcheck import (
+    normalize,
     training_args,
 )
 
@@ -15,3 +27,86 @@ def test_ema_checkpoint_retention_is_left_for_runtime_inheritance() -> None:
     training_argument.check_value(normalized, strict=True)
 
     assert normalized["ema_ckpt_keep"] is None
+
+
+@pytest.mark.parametrize("reset", [None, False, True])
+def test_reset_data_progress_is_an_opt_in_boolean(reset) -> None:
+    training_argument = training_args()
+    values = {"numb_steps": 1}
+    if reset is not None:
+        values["reset_data_progress"] = reset
+
+    normalized = training_argument.normalize_value(values)
+    training_argument.check_value(normalized, strict=True)
+
+    assert normalized["reset_data_progress"] is (False if reset is None else reset)
+
+
+@pytest.mark.parametrize("reset", [None, 0, 1, "true", "false", [], {}])
+def test_reset_data_progress_rejects_non_boolean_values(reset) -> None:
+    training_argument = training_args()
+    normalized = training_argument.normalize_value(
+        {"numb_steps": 1, "reset_data_progress": reset}
+    )
+
+    with pytest.raises(ArgumentTypeError, match="reset_data_progress"):
+        training_argument.check_value(normalized, strict=True)
+
+
+def _training_input(data: dict[str, Any] | list[dict[str, Any]]) -> dict[str, Any]:
+    return {
+        "model": {
+            "type_map": ["H"],
+            "descriptor": {"type": "se_e2_a", "sel": [1]},
+            "fitting_net": {},
+        },
+        "training": {"numb_steps": 1, "training_data": data},
+    }
+
+
+def test_training_groups_normalize_aliases_without_mutating_input() -> None:
+    source = _training_input(
+        [
+            {
+                "systems": "first.lmdb",
+                "prob": 0.6,
+                "auto_prob_style": "prob_sys_size;0:1:1",
+            },
+            {"systems": ["second.lmdb", "third.lmdb"], "prob": 0.4},
+        ]
+    )
+    original = copy.deepcopy(source)
+
+    groups = normalize(source)["training"]["training_data"]
+
+    assert source == original
+    assert groups[0]["auto_prob"] == "prob_sys_size;0:1:1"
+    assert "auto_prob_style" not in groups[0]
+    assert groups[1]["auto_prob"] == "prob_sys_size"
+    assert groups[0]["batch_size"] == groups[1]["batch_size"] == "auto"
+    assert [group["prob"] for group in groups] == [0.6, 0.4]
+
+
+def test_legacy_training_dictionary_remains_a_dictionary() -> None:
+    group = normalize(_training_input({"systems": "dataset"}))["training"][
+        "training_data"
+    ]
+
+    assert isinstance(group, dict)
+    assert group["auto_prob"] == "prob_sys_size"
+    assert group["batch_size"] == "auto"
+
+
+@pytest.mark.parametrize(
+    "groups",
+    [
+        [],
+        [{"systems": "dataset", "unknown_option": True}],
+        [{"systems": "dataset", "prob": -0.1}],
+        [{"systems": "dataset", "prob": 0.6}],
+        [{"systems": "dataset", "prob": float("nan")}],
+    ],
+)
+def test_invalid_training_groups_are_rejected(groups: list[dict[str, Any]]) -> None:
+    with pytest.raises((ArgumentError, ValueError)):
+        normalize(_training_input(groups))
diff --git a/source/tests/common/test_auto_batch_size.py b/source/tests/common/test_auto_batch_size.py
index 16ec6fbef7..3dfd44568a 100644
--- a/source/tests/common/test_auto_batch_size.py
+++ b/source/tests/common/test_auto_batch_size.py
@@ -170,3 +170,44 @@ def evaluate(batch: Any) -> Any:
         with self.assertRaises(OutOfMemoryError):
             auto.execute_all(evaluate, 1, 2, xp.zeros((1, 2)))
         self.assertEqual(attempts, [1])
+
+
+class TestLargeFrameBatchGrowth(unittest.TestCase):
+    @patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "0"})
+    def test_large_frames_grow_from_successful_work(self) -> None:
+        auto = CustomizedAutoBatchSizeGPU(1024)
+        attempts = []
+
+        def evaluate(nframes: int, start: int) -> tuple[int, None]:
+            attempts.append(nframes)
+            return nframes, None
+
+        for _ in range(3):
+            auto.execute(evaluate, 0, 10000)
+        self.assertEqual(attempts, [1, 2, 4])
+        self.assertEqual(auto.current_batch_size, 80000)
+
+    def test_large_frames_do_not_expand_explicit_or_cpu_limits(self) -> None:
+        for cls, limit in (
+            (CustomizedAutoBatchSizeCPU, "0"),
+            (CustomizedAutoBatchSizeGPU, "1024"),
+        ):
+            with (
+                self.subTest(cls=cls, limit=limit),
+                patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": limit}),
+            ):
+                auto = cls(1024)
+                for _ in range(3):
+                    nframes, _ = auto.execute(lambda n, start: (n, None), 0, 10000)
+                    self.assertEqual(nframes, 1)
+                self.assertEqual(auto.current_batch_size, 1024)
+
+    @patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "0"})
+    def test_partial_tail_does_not_shrink_learned_budget(self) -> None:
+        auto = CustomizedAutoBatchSizeGPU(1024)
+        for natoms in (100, 101, 200):
+            nframes, _ = auto.execute(
+                lambda n, start: (n, None), 0, natoms, max_nframes=2
+            )
+            self.assertEqual(nframes, 2)
+            self.assertEqual(auto.current_batch_size, 1024)
diff --git a/source/tests/common/test_data_config.py b/source/tests/common/test_data_config.py
new file mode 100644
index 0000000000..01a5b11481
--- /dev/null
+++ b/source/tests/common/test_data_config.py
@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Shared dataset configuration and sampling contracts."""
+
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+import pytest
+
+from deepmd.utils import random as dp_random
+from deepmd.utils.data_config import (
+    resolve_data_groups,
+)
+from deepmd.utils.data_sampling import (
+    group_probabilities,
+    system_probabilities,
+)
+from deepmd.utils.data_system import (
+    get_data,
+)
+
+
+def test_group_probabilities_do_not_equalize_unweighted_groups() -> None:
+    np.testing.assert_allclose(group_probabilities([2, 8], [None, None]), [0.2, 0.8])
+    np.testing.assert_allclose(group_probabilities([2, 8], [0.6, 0.4]), [0.6, 0.4])
+    np.testing.assert_allclose(
+        group_probabilities([2, 3, 9], [0.6, None, None]), [0.6, 0.1, 0.3]
+    )
+    np.testing.assert_allclose(group_probabilities([2, 8], [1.0, None]), [1.0, 0.0])
+
+
+def test_group_and_system_probabilities_are_independent() -> None:
+    outer = group_probabilities([10, 90], [0.6, 0.4])
+    inner = system_probabilities([1, 2, 7], auto_prob="prob_sys_size;0:1:0.5;1:3:0.5")
+
+    np.testing.assert_allclose(outer[0] * inner, [0.3, 0.6 / 9, 2.1 / 9])
+    assert np.isclose((outer[0] * inner).sum(), 0.6)
+
+
+@pytest.mark.parametrize("systems_as_list", [False, True])
+def test_multi_lmdb_accepts_size_and_uniform_auto_prob(
+    tmp_path: Path, systems_as_list: bool
+) -> None:
+    first = tmp_path / "first.lmdb"
+    second = tmp_path / "nested" / "second.lmdb"
+    first.mkdir()
+    second.mkdir(parents=True)
+    systems = [str(first), str(second)] if systems_as_list else str(tmp_path)
+
+    groups = resolve_data_groups({"systems": systems})
+    assert set(groups[0].systems) == {str(first), str(second)}
+    assert groups[0].auto_prob is None
+    explicit_default = resolve_data_groups(
+        {"systems": systems, "auto_prob": "prob_sys_size"}
+    )
+    assert explicit_default[0].systems == groups[0].systems
+    uniform = resolve_data_groups({"systems": systems, "auto_prob": "prob_uniform"})
+    assert uniform[0].systems == groups[0].systems
+    with pytest.raises(ValueError, match="auto_prob must be prob_sys_size"):
+        resolve_data_groups({"systems": systems, "auto_prob": "prob_sys_size;0:1:1"})
+    with pytest.raises(ValueError, match="sys_probs is not supported"):
+        resolve_data_groups({"systems": systems, "sys_probs": [-1, -1]})
+
+
+def test_empty_groups_cannot_receive_positive_probability() -> None:
+    with pytest.raises(ValueError, match="must contain batches"):
+        group_probabilities([0, 8], [0.5, 0.5])
+    np.testing.assert_allclose(group_probabilities([0, 8], [None, None]), [0, 1])
+
+
+def test_single_lmdb_groups_retain_internal_auto_prob() -> None:
+    groups = resolve_data_groups(
+        [
+            {
+                "systems": "first.lmdb",
+                "auto_prob": "prob_sys_size;0:10:0.5;10:30:0.5",
+                "prob": 0.6,
+            },
+            {"systems": "second.lmdb", "prob": 0.4},
+        ]
+    )
+
+    assert groups[0].auto_prob == "prob_sys_size;0:10:0.5;10:30:0.5"
+    assert groups[1].auto_prob is None
+    assert [group.prob for group in groups] == [0.6, 0.4]
+
+
+def test_numpy_groups_follow_the_training_seed() -> None:
+    system = str(Path(__file__).resolve().parents[1] / "pt" / "property" / "double")
+    config = [{"systems": system, "batch_size": size} for size in (1, 2)]
+    sequences = []
+    for _ in range(2):
+        dp_random.seed(123)
+        data = get_data(config, 0, None, None)
+        sequences.append([data.get_batch()["coord"].shape[0] for _ in range(24)])
+    assert sequences[0] == sequences[1]
+    assert set(sequences[0]) == {1, 2}
+
+
+def test_numpy_backend_parent_discovery_preserves_format_scope(tmp_path: Path) -> None:
+    numpy = tmp_path / "numpy"
+    numpy.mkdir()
+    (numpy / "type.raw").touch()
+    lmdb = tmp_path / "frames.lmdb"
+    lmdb.mkdir()
+    groups = resolve_data_groups({"systems": str(tmp_path)}, formats=("deepmd",))
+    assert groups[0].systems == (str(numpy),)
diff --git a/source/tests/common/test_deepmd_data.py b/source/tests/common/test_deepmd_data.py
index 9e9e57d9eb..976e2302f9 100644
--- a/source/tests/common/test_deepmd_data.py
+++ b/source/tests/common/test_deepmd_data.py
@@ -8,7 +8,9 @@
     Path,
 )
 
+import h5py
 import numpy as np
+import pytest
 
 from deepmd.dpmodel.utils.dist_check import (
     compute_min_pair_margin_single,
@@ -22,6 +24,57 @@
 )
 
 
+def test_cached_numpy_set_applies_copy_weights_once(tmp_path: Path) -> None:
+    np.savetxt(tmp_path / "type.raw", [0], fmt="%d")
+    set_path = tmp_path / "set.000"
+    set_path.mkdir()
+    np.save(set_path / "coord.npy", [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
+    np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (2, 1)))
+    np.save(set_path / "numb_copy.npy", [2, 1])
+    data = DeepmdData(str(tmp_path))
+
+    assert data.get_mean_batch_nframes(8) == 3
+    for _ in range(3):
+        batch = data.get_batch(8)
+        np.testing.assert_array_equal(np.sort(batch["coord"][:, 0]), [0, 0, 1])
+
+
+@pytest.mark.parametrize("hdf5", [False, True])
+def test_uniform_numpy_counts_undersized_sets(tmp_path: Path, hdf5: bool) -> None:
+    paths = []
+    container = tmp_path / "systems.h5"
+    for index, sizes in enumerate(((1, 5), (2, 8))):
+        root = tmp_path / f"system_{index}"
+        root.mkdir()
+        np.savetxt(root / "type.raw", [0], fmt="%d")
+        for set_index, size in enumerate(sizes):
+            set_path = root / f"set.{set_index:03d}"
+            set_path.mkdir()
+            np.save(set_path / "coord.npy", np.full((size, 3), index, dtype=np.float64))
+            np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (size, 1)))
+        paths.append(str(root))
+        if hdf5:
+            with h5py.File(container, "a") as handle:
+                group = handle.create_group(root.name)
+                group.create_dataset("type.raw", data=[0])
+                for file in root.rglob("*.npy"):
+                    group.create_dataset(
+                        str(file.relative_to(root)), data=np.load(file)
+                    )
+            paths[-1] = f"{container}#/{root.name}"
+    data = DeepmdDataSystem(paths, [3, 5], 1, auto_prob_style="prob_uniform")
+    means = [
+        np.mean([data.get_batch(sys_idx=index)["coord"].shape[0] for _ in range(4)])
+        for index in range(2)
+    ]
+    np.testing.assert_allclose(means, [2, 3.5])
+    np.testing.assert_allclose(data.sys_probs * means, [14 / 11, 14 / 11])
+    data = DeepmdDataSystem(paths, "mixed:4", 1, auto_prob_style="prob_uniform")
+    np.testing.assert_allclose(data.sys_probs, [0.5, 0.5])
+    assert data.uniform_sampling_mass == 0.5
+    assert data.get_batch()["coord"].shape[0] == 4
+
+
 class TestDeepmdDataTypeMap(unittest.TestCase):
     def setUp(self) -> None:
         self.tmpdir = tempfile.TemporaryDirectory()
diff --git a/source/tests/common/test_deepmd_data_grid_density.py b/source/tests/common/test_deepmd_data_grid_density.py
new file mode 100644
index 0000000000..2af6404cab
--- /dev/null
+++ b/source/tests/common/test_deepmd_data_grid_density.py
@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Tests for the frame-major grid/density branches in DeepmdData.
+
+These keys are loaded verbatim (no reshaping to natoms) and must carry a
+leading frame dimension matching the set's frame count, so that shuffling
+never pairs a structure with another frame's grid or density label.
+"""
+
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+
+from deepmd.utils.data import (
+    DeepmdData,
+)
+
+NATOMS = 3
+NGRID = 5
+NFRAMES = 4
+
+
+class TestDeepmdDataGridDensity(unittest.TestCase):
+    def setUp(self) -> None:
+        self.tmpdir = tempfile.TemporaryDirectory()
+        self.root = Path(self.tmpdir.name)
+        self.set_dir = self.root / "set.000"
+        self.set_dir.mkdir()
+        rng = np.random.default_rng(42)
+
+        np.savetxt(self.root / "type.raw", np.zeros(NATOMS, dtype=np.int32), fmt="%d")
+        np.save(
+            self.set_dir / "coord.npy",
+            rng.random((NFRAMES, NATOMS * 3), dtype=np.float32),
+        )
+        np.save(
+            self.set_dir / "box.npy",
+            np.eye(3, dtype=np.float32).reshape(1, 9).repeat(NFRAMES, axis=0),
+        )
+        np.save(
+            self.set_dir / "type.npy",
+            np.zeros((NFRAMES, NATOMS), dtype=np.int32),
+        )
+        self.set_data("grid", rng.random((NFRAMES, NGRID, 3), dtype=np.float32))
+        self.set_data("density", rng.random((NFRAMES, NGRID, 1), dtype=np.float32))
+
+    def tearDown(self) -> None:
+        self.tmpdir.cleanup()
+
+    def set_data(self, key: str, arr: np.ndarray | None) -> None:
+        path = self.set_dir / f"{key}.npy"
+        if arr is None:
+            path.unlink(missing_ok=True)
+        else:
+            np.save(path, arr)
+
+    def build_data(self, must: bool = True) -> DeepmdData:
+        data = DeepmdData(str(self.root))
+        data.add("coord", 3, atomic=True, must=True)
+        data.add("box", 9, atomic=False, must=True)
+        data.add("type", 1, atomic=True, must=True)
+        data.add(
+            "grid",
+            3,
+            atomic=False,
+            must=must,
+            high_prec=True,
+            special_shape="frame_major",
+        )
+        data.add(
+            "density",
+            1,
+            atomic=False,
+            must=must,
+            high_prec=True,
+            special_shape="frame_major",
+        )
+        return data
+
+    def test_load_set_frame_major(self) -> None:
+        data = self.build_data()
+        loaded = data._load_set(self.set_dir)
+        self.assertEqual(loaded["find_grid"], 1.0)
+        self.assertEqual(loaded["find_density"], 1.0)
+        self.assertEqual(loaded["grid"].shape, (NFRAMES, NGRID, 3))
+        self.assertEqual(loaded["density"].shape, (NFRAMES, NGRID, 1))
+
+    def test_iter_test_uses_frame_major_extent_for_chunking(self) -> None:
+        data = self.build_data()
+        chunks = list(data.iter_test(chunk_atoms=NGRID * 2))
+        self.assertEqual([chunk["box"].shape[0] for chunk in chunks], [2, 2])
+
+    def test_single_frame(self) -> None:
+        data = self.build_data()
+        frame = data.get_single_frame(0, 0)
+        self.assertEqual(frame["find_grid"], 1.0)
+        self.assertEqual(frame["find_density"], 1.0)
+        self.assertEqual(frame["grid"].shape, (NGRID, 3))
+        self.assertEqual(frame["density"].shape, (NGRID, 1))
+
+    def test_frame_count_mismatch(self) -> None:
+        self.set_data("grid", np.zeros((NFRAMES - 1, NGRID, 3), dtype=np.float32))
+        loaders = {
+            "set": lambda data: data._load_set(self.set_dir),
+            "single frame": lambda data: data.get_single_frame(0, 0),
+        }
+        for name, load in loaders.items():
+            with self.subTest(loader=name):
+                with self.assertRaisesRegex(ValueError, "frame count"):
+                    load(self.build_data())
+
+    def test_missing_frame_dimension(self) -> None:
+        # stored without a leading frame dimension: rejected rather than
+        # silently misaligned with the frames
+        self.set_data("grid", np.zeros((NGRID, 3), dtype=np.float32))
+        loaders = {
+            "set": lambda data: data._load_set(self.set_dir),
+            "single frame": lambda data: data.get_single_frame(0, 0),
+        }
+        for name, load in loaders.items():
+            with self.subTest(loader=name):
+                with self.assertRaisesRegex(ValueError, r"frame count|leading frame"):
+                    load(self.build_data())
+
+    def test_missing_file_optional(self) -> None:
+        # optional labels absent: the loader must not fail and must signal
+        # find=0 so that downstream code can skip the residual
+        self.set_data("density", None)
+        data = self.build_data(must=False)
+        loaded = data._load_set(self.set_dir)
+        self.assertEqual(loaded["find_density"], 0.0)
+        frame = data.get_single_frame(0, 0)
+        self.assertEqual(frame["find_density"], 0.0)
+
+    def test_invalid_frame_major_layout(self) -> None:
+        invalid = (
+            ("wrong trailing width", (NFRAMES, NGRID, 4), "ndof"),
+            ("divisible wrong width", (NFRAMES, NGRID, 6), "ndof"),
+            ("invalid flattened width", (NFRAMES, NGRID * 3 + 1), "ndof"),
+            ("empty", (NFRAMES, 0, 3), "at least one point"),
+            ("rank four", (NFRAMES, 1, NGRID, 3), "must have shape"),
+        )
+        loaders = {
+            "set": lambda data: data._load_set(self.set_dir),
+            "single frame": lambda data: data.get_single_frame(0, 0),
+        }
+        for case, shape, message in invalid:
+            self.set_data("grid", np.zeros(shape, dtype=np.float32))
+            for loader_name, load in loaders.items():
+                with self.subTest(case=case, loader=loader_name):
+                    with self.assertRaisesRegex(ValueError, message):
+                        load(self.build_data())
+
+        # A correctly flattened frame-major array remains accepted.
+        self.set_data("grid", np.zeros((NFRAMES, NGRID * 3), dtype=np.float32))
+        loaded = self.build_data()._load_set(self.set_dir)
+        self.assertEqual(loaded["find_grid"], 1.0)
+
+    def test_ragged_mixed_batch_clear_error(self) -> None:
+        # mixed batching across systems with different ngrid must fail with
+        # a clear message instead of a raw numpy shape error
+        from deepmd.utils.data_system import (
+            DeepmdDataSystem,
+        )
+
+        class StubDS(DeepmdDataSystem):
+            """Test double: instantiated via __new__ to skip the heavy init."""
+
+            def get_ntypes(self) -> int:
+                return 2
+
+            def get_data_dict(self, idx: int) -> dict:
+                return {
+                    "density": {
+                        "atomic": False,
+                        "special_shape": "frame_major",
+                        "ndof": 1,
+                    }
+                }
+
+        batch_data = [
+            {
+                "type": np.zeros((1, 3), dtype=np.int32),
+                "natoms_vec": np.array([3, 3, 1, 1], dtype=np.int32),
+                "default_mesh": np.zeros(9),
+                "density": np.zeros((1, 5), dtype=np.float32),
+                "find_density": 1.0,
+            },
+            {
+                "type": np.zeros((1, 3), dtype=np.int32),
+                "natoms_vec": np.array([3, 3, 1, 1], dtype=np.int32),
+                "default_mesh": np.zeros(9),
+                "density": np.zeros((1, 6), dtype=np.float32),
+                "find_density": 1.0,
+            },
+        ]
+        with self.assertRaisesRegex(ValueError, "extent"):
+            StubDS.__new__(StubDS)._merge_batch_data(batch_data)
+
+    def test_grid_density_extent_mismatch(self) -> None:
+        # grid and density must carry the same number of points per frame
+        self.set_data("density", np.zeros((NFRAMES, NGRID + 2, 1), dtype=np.float32))
+        with self.assertRaisesRegex(ValueError, "extents"):
+            self.build_data()._load_set(self.set_dir)
+
+    def test_extent_mismatch_single_frame(self) -> None:
+        # the PT training path (get_item_torch -> get_single_frame) assembles
+        # each frame key by key; the extent cross-check must run there too,
+        # or a 1-point density broadcasts silently against the ngrid-point
+        # prediction in the loss
+        self.set_data("density", np.zeros((NFRAMES, 1, 1), dtype=np.float32))
+        data = self.build_data()
+        with self.assertRaisesRegex(ValueError, "extents"):
+            data.get_item_torch(0)
+
+    def test_cross_set_extent_mismatch(self) -> None:
+        # two sets of one system must agree on the frame-major extent;
+        # numpy would otherwise fail with a raw shape error in get_test
+        set1 = self.root / "set.001"
+        set1.mkdir()
+        rng = np.random.default_rng(1)
+        np.save(set1 / "coord.npy", rng.random((NFRAMES, NATOMS * 3), dtype=np.float32))
+        np.save(
+            set1 / "box.npy",
+            np.eye(3, dtype=np.float32).reshape(1, 9).repeat(NFRAMES, axis=0),
+        )
+        np.save(set1 / "type.npy", np.zeros((NFRAMES, NATOMS), dtype=np.int32))
+        np.save(
+            set1 / "grid.npy", rng.random((NFRAMES, NGRID + 3, 3), dtype=np.float32)
+        )
+        np.save(
+            set1 / "density.npy", rng.random((NFRAMES, NGRID + 3, 1), dtype=np.float32)
+        )
+        with self.assertRaisesRegex(ValueError, "extents"):
+            self.build_data().get_test()
+
+    def test_property_named_density_not_hijacked(self) -> None:
+        # a user property named "density" without a special_shape declaration
+        # must load through the ordinary path, not the frame-major branch
+        self.set_data("density", np.zeros((NFRAMES, 2), dtype=np.float32))
+        data = DeepmdData(str(self.root))
+        data.add("coord", 3, atomic=True, must=True)
+        data.add("box", 9, atomic=False, must=True)
+        data.add("type", 1, atomic=True, must=True)
+        data.add("density", 2, atomic=False, must=True)
+        loaded = data._load_set(self.set_dir)
+        self.assertEqual(loaded["find_density"], 1.0)
+        self.assertEqual(loaded["density"].shape, (NFRAMES, 2))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/common/test_directory_data_progress.py b/source/tests/common/test_directory_data_progress.py
new file mode 100644
index 0000000000..a57ae77fca
--- /dev/null
+++ b/source/tests/common/test_directory_data_progress.py
@@ -0,0 +1,384 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Compact, lazy batch-progress restoration for directory and HDF5 sources."""
+
+import copy
+import json
+from unittest.mock import (
+    patch,
+)
+
+import h5py
+import numpy as np
+import pytest
+
+from deepmd.utils import random as dp_random
+from deepmd.utils.data import (
+    DeepmdData,
+)
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
+
+
+def _write_system(root, *, hdf5=False, nsets=1, copies=False, offset=0, nframes=7):
+    """Write uniquely identifiable frames with optional zero/repeated copies."""
+    arrays = {"type.raw": np.array([0, 0], dtype=np.int32)}
+    for index in range(nsets):
+        size = nframes + index
+        frame_ids = offset + index * 100 + np.arange(size)
+        prefix = f"set.{index:03d}/"
+        arrays[prefix + "coord.npy"] = np.repeat(frame_ids[:, None], 6, axis=1)
+        arrays[prefix + "box.npy"] = np.tile(np.eye(3).reshape(1, 9), (size, 1))
+        arrays[prefix + "energy.npy"] = frame_ids[:, None] * 10
+        if copies:
+            arrays[prefix + "numb_copy.npy"] = (np.arange(size) % 4)[:, None]
+    if hdf5:
+        filename = root.with_suffix(".h5")
+        with h5py.File(filename, "w") as handle:
+            for name, array in arrays.items():
+                handle.create_dataset("system/" + name, data=array)
+        return str(filename) + "#/system"
+    root.mkdir()
+    for name, array in arrays.items():
+        target = root / name
+        target.parent.mkdir(exist_ok=True)
+        if name.endswith(".raw"):
+            np.savetxt(target, array, fmt="%d")
+        else:
+            np.save(target, array)
+    return str(root)
+
+
+def _assert_batch_equal(actual, expected):
+    assert actual.keys() == expected.keys()
+    for name in actual:
+        np.testing.assert_array_equal(actual[name], expected[name], err_msg=name)
+
+
+@pytest.mark.parametrize("hdf5", [False, True], ids=["directory", "hdf5"])
+@pytest.mark.parametrize("nsets", [1, 3])
+@pytest.mark.parametrize("copies", [False, True])
+@pytest.mark.parametrize("consumed", [0, 1, 3, 11])
+def test_reader_resume_multiple_rounds(tmp_path, hdf5, nsets, copies, consumed):
+    """Resume before/within/after sets, including previously reshuffled caches."""
+    path = _write_system(tmp_path / "data", hdf5=hdf5, nsets=nsets, copies=copies)
+    original = DeepmdData(path)
+    original.add("energy", 1)
+    dp_random.seed(482)
+    for _ in range(consumed):
+        original.get_batch(2)
+    progress = original.state_dict()
+    rng_state = dp_random.get_state()
+    expected = [(original.get_batch(2), dp_random.random()) for _ in range(40)]
+
+    restored = DeepmdData(path)
+    restored.load_state_dict(progress)
+    # Requirements can be installed after logical state restoration.
+    restored.add("energy", 1)
+    dp_random.set_state(rng_state)
+    for batch, random_value in expected:
+        _assert_batch_equal(restored.get_batch(2), batch)
+        assert dp_random.random() == random_value
+
+
+@pytest.mark.parametrize("hdf5", [False, True], ids=["directory", "hdf5"])
+def test_restore_is_lazy_and_invalidates_changed_requirements(tmp_path, hdf5):
+    path = _write_system(tmp_path / "data", hdf5=hdf5)
+    data = DeepmdData(path)
+    data.get_batch(2)
+    progress = data.state_dict()
+    data.get_test()
+    assert hasattr(data, "batch_set")
+    assert hasattr(data, "test_set")
+    data.add("energy", 1)
+    with (
+        patch.object(data, "_load_set", side_effect=AssertionError("decoded set")),
+        patch.object(data, "_get_nframes", side_effect=AssertionError("scanned set")),
+        patch.object(
+            type(data.dirs[0]), "load_numpy", side_effect=AssertionError("read data")
+        ),
+    ):
+        data.load_state_dict(progress)
+        assert not hasattr(data, "batch_set")
+        assert not hasattr(data, "test_set")
+    with patch.object(data, "_load_set", wraps=data._load_set) as load:
+        batch = data.get_batch(2)
+        load.assert_called_once()
+    np.testing.assert_array_equal(batch["energy"][:, 0], batch["coord"][:, 0] * 10)
+
+
+@pytest.mark.parametrize("extra_batches", [1, 5])
+def test_in_process_restore_reuses_canonical_cache(tmp_path, extra_batches):
+    """Display/statistics snapshots must not force another set decode."""
+    path = _write_system(tmp_path / "data")
+    data = DeepmdData(path)
+    data.get_batch(2)
+    state = data.state_dict()
+    expected = data.get_batch(2)
+    for _ in range(extra_batches):
+        data.get_batch(2)
+    cached_set = data.batch_set
+    with patch.object(data, "_load_set", side_effect=AssertionError("decoded set")):
+        data.load_state_dict(state)
+        _assert_batch_equal(data.get_batch(2), expected)
+    assert data.batch_set is cached_set
+
+
+def test_in_process_restore_discards_other_set_cache(tmp_path):
+    path = _write_system(tmp_path / "data", nsets=2)
+    data = DeepmdData(path)
+    data.get_batch(2)
+    state = data.state_dict()
+    expected = data.get_batch(2)
+    for _ in range(2):
+        data.get_batch(2)
+    assert data._batch_set_name == data.dirs[1]
+    data.load_state_dict(state)
+    assert not hasattr(data, "batch_set")
+    _assert_batch_equal(data.get_batch(2), expected)
+
+
+def test_exhausted_set_is_not_decoded_on_resume(tmp_path):
+    path = _write_system(tmp_path / "data", nsets=2)
+    original = DeepmdData(path)
+    for _ in range(3):
+        original.get_batch(2)
+    state = original.state_dict()
+    restored = DeepmdData(path)
+    restored.load_state_dict(state)
+    with patch.object(restored, "_load_set", wraps=restored._load_set) as load:
+        restored.get_batch(2)
+    assert [call.args[0] for call in load.call_args_list] == [restored.dirs[1]]
+
+
+def test_numb_copy_does_not_compound_between_rounds(tmp_path):
+    path = _write_system(tmp_path / "data", copies=True)
+    data = DeepmdData(path)
+    frame_ids = np.arange(7)
+    expected_ids = np.repeat(frame_ids, frame_ids % 4)
+    for _ in range(5):
+        batch = data.get_batch(len(expected_ids))
+        np.testing.assert_array_equal(np.sort(batch["coord"][:, 0]), expected_ids)
+        assert data.batch_set["coord"].shape[0] == 7
+
+
+@pytest.mark.parametrize("copies", [False, True])
+def test_batch_order_matches_canonical_shuffle(tmp_path, copies):
+    """Index-only batching retains the established shuffle/RNG semantics."""
+    path = _write_system(tmp_path / "data", nsets=3, copies=copies)
+    data = DeepmdData(path)
+    dp_random.seed(716)
+    for index in range(6):
+        rng_state = dp_random.get_state()
+        expected, _ = data._shuffle_data(data._load_set(data.dirs[index % 3]))
+        expected_rng = dp_random.get_state()
+        dp_random.set_state(rng_state)
+        # An oversized batch returns the entire set, preserving legacy
+        # short-batch behavior rather than crossing into the following set.
+        _assert_batch_equal(data.get_batch(1000), expected)
+        assert dp_random.get_state() == expected_rng
+
+
+@pytest.mark.parametrize("hdf5", [False, True], ids=["directory", "hdf5"])
+@pytest.mark.parametrize("batch_size", [2, "mixed:3"])
+def test_system_resume_random_choices(tmp_path, hdf5, batch_size):
+    paths = [
+        _write_system(
+            tmp_path / f"data{index}",
+            hdf5=hdf5,
+            nsets=index + 1,
+            copies=True,
+            offset=index * 1000,
+        )
+        for index in range(2)
+    ]
+    original = DeepmdDataSystem(paths, batch_size, 1, sys_probs=[0.3, 0.7])
+    dp_random.seed(918)
+    for _ in range(13):
+        original.get_batch()
+    progress = original.state_dict()
+    rng_state = dp_random.get_state()
+    expected = [original.get_batch() for _ in range(60)]
+    restored = DeepmdDataSystem(paths, batch_size, 1, sys_probs=[0.3, 0.7])
+    with patch.object(DeepmdData, "_load_set", side_effect=AssertionError("decoded")):
+        restored.load_state_dict(progress)
+    dp_random.set_state(rng_state)
+    for batch in expected:
+        _assert_batch_equal(restored.get_batch(), batch)
+
+
+@pytest.mark.parametrize("batch_size", [2, "mixed:3"])
+def test_system_resume_restores_exact_probabilities_after_roundoff(
+    tmp_path, batch_size
+):
+    paths = [
+        _write_system(tmp_path / f"data{index}", offset=index * 1000)
+        for index in range(2)
+    ]
+    # Put the first random draw exactly on a selection threshold. Even a
+    # one-ULP probability difference then changes the selected system.
+    dp_random.seed(1)
+    threshold = float(dp_random.random())
+    saved_probs = np.array([threshold, 1 - threshold])
+    rounded_threshold = np.nextafter(threshold, np.inf)
+    recomputed_probs = np.array([rounded_threshold, 1 - rounded_threshold])
+    dp_random.seed(1)
+    assert dp_random.choice(2, p=saved_probs) == 1
+    dp_random.seed(1)
+    assert dp_random.choice(2, p=recomputed_probs) == 0
+
+    original = DeepmdDataSystem(paths, batch_size, 1, sys_probs=saved_probs.tolist())
+    state = original.state_dict()
+    dp_random.seed(1)
+    rng_state = dp_random.get_state()
+    expected = [original.get_batch() for _ in range(20)]
+    expected_rng_state = dp_random.get_state()
+
+    restored = DeepmdDataSystem(
+        paths, batch_size, 1, sys_probs=recomputed_probs.tolist()
+    )
+    with patch.object(DeepmdData, "_load_set", side_effect=AssertionError("decoded")):
+        restored.load_state_dict(state)
+    np.testing.assert_array_equal(restored.sys_probs, saved_probs)
+    # The restored probabilities must not alias the caller's checkpoint.
+    state["sys_probs"][0] = 0
+    np.testing.assert_array_equal(restored.sys_probs, saved_probs)
+    dp_random.set_state(rng_state)
+    for batch in expected:
+        _assert_batch_equal(restored.get_batch(), batch)
+    assert dp_random.get_state() == expected_rng_state
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_system_resume_preserves_exactly_matching_probability_dtype(tmp_path, dtype):
+    paths = [
+        _write_system(tmp_path / f"data{index}", offset=index * 1000)
+        for index in range(3)
+    ]
+    # Float32 thirds are valid for NumPy choice but their float64 sum is
+    # outside its tighter float64 normalization tolerance.
+    probabilities = list(np.full(3, 1 / 3, dtype=dtype))
+    original = DeepmdDataSystem(paths, 2, 1, sys_probs=probabilities)
+    state = original.state_dict()
+    dp_random.seed(193)
+    rng_state = dp_random.get_state()
+    expected = [original.get_batch() for _ in range(20)]
+
+    restored = DeepmdDataSystem(paths, 2, 1, sys_probs=probabilities)
+    restored.load_state_dict(state)
+    assert restored.sys_probs.dtype == dtype
+    np.testing.assert_array_equal(restored.sys_probs, probabilities)
+    dp_random.set_state(rng_state)
+    for batch in expected:
+        _assert_batch_equal(restored.get_batch(), batch)
+
+
+@pytest.mark.parametrize(
+    "saved_probs",
+    [
+        None,
+        0.5,
+        [0.5],
+        [[0.3, 0.7]],
+        [0.3, 0.7, 0.0],
+        [float("nan"), 0.7],
+        [float("inf"), 0.7],
+        [-0.3, 1.3],
+        ["invalid", 0.7],
+        [0.4, 0.6],
+        [0.3 + 1e-8, 0.7 - 1e-8],
+    ],
+)
+def test_invalid_system_probabilities_are_rejected_without_mutation(
+    tmp_path, saved_probs
+):
+    paths = [_write_system(tmp_path / f"data{index}") for index in range(2)]
+    data = DeepmdDataSystem(paths, 2, 1, sys_probs=[0.3, 0.7])
+    data.get_batch()
+    initial = data.state_dict()
+    invalid = copy.deepcopy(initial)
+    invalid["sys_probs"] = saved_probs
+    with pytest.raises(ValueError, match="probabilities"):
+        data.load_state_dict(invalid)
+    assert data.state_dict() == initial
+
+
+@pytest.mark.parametrize("saved_zero", [False, True])
+def test_system_probability_roundoff_does_not_change_excluded_systems(
+    tmp_path, saved_zero
+):
+    paths = [_write_system(tmp_path / f"data{index}") for index in range(2)]
+    positive = np.nextafter(0.0, 1.0)
+    current_probs = [positive if saved_zero else 0.0, 1.0]
+    data = DeepmdDataSystem(paths, 2, 1, sys_probs=current_probs)
+    initial = data.state_dict()
+    invalid = copy.deepcopy(initial)
+    invalid["sys_probs"] = [0.0 if saved_zero else positive, 1.0]
+    with pytest.raises(ValueError, match="probabilities"):
+        data.load_state_dict(invalid)
+    assert data.state_dict() == initial
+
+
+def test_system_probabilities_are_not_restored_before_source_validation(tmp_path):
+    paths = [_write_system(tmp_path / f"data{index}") for index in range(2)]
+    data = DeepmdDataSystem(paths, 2, 1, sys_probs=[0.3, 0.7])
+    initial = data.state_dict()
+    invalid = copy.deepcopy(initial)
+    invalid["sys_probs"][0] = np.nextafter(0.3, 1.0)
+    invalid["systems"][-1]["version"] = 99
+    with pytest.raises(ValueError):
+        data.load_state_dict(invalid)
+    assert data.state_dict() == initial
+
+
+def test_checkpoint_is_compact_and_detached(tmp_path):
+    path = _write_system(tmp_path / "data", nframes=10000)
+    data = DeepmdData(path)
+    data.get_batch(2)
+    state = data.state_dict()
+    # The fixed-size RNG state is JSON-serializable; no frame permutation or
+    # label arrays should scale this snapshot with the number of frames.
+    assert len(json.dumps(state)) < 12000
+    assert state["set_size"] == 10000
+    state["identity"]["prefix_sum"][0] = -1
+    state["shuffle_state"][1][0] = 0
+    assert data.state_dict()["identity"]["prefix_sum"][0] == 10000
+    assert data.state_dict()["shuffle_state"][1][0] != 0
+
+
+@pytest.mark.parametrize("change", ["source", "batch_size", "sys_probs", "order"])
+def test_incompatible_system_progress_is_rejected(tmp_path, change):
+    paths = [_write_system(tmp_path / f"data{index}") for index in range(2)]
+    original = DeepmdDataSystem(paths, 2, 1, sys_probs=[0.3, 0.7])
+    original.get_batch()
+    state = original.state_dict()
+    kwargs = {"batch_size": 2, "test_size": 1, "sys_probs": [0.3, 0.7]}
+    if change == "source":
+        paths[0] = _write_system(tmp_path / "other")
+    elif change == "batch_size":
+        kwargs["batch_size"] = 1
+    elif change == "sys_probs":
+        kwargs["sys_probs"] = [0.4, 0.6]
+    elif change == "order":
+        paths.reverse()
+    restored = DeepmdDataSystem(paths, **kwargs)
+    initial = restored.state_dict()
+    with pytest.raises(ValueError):
+        restored.load_state_dict(state)
+    assert restored.state_dict() == initial
+
+
+@pytest.mark.parametrize(
+    "field,value",
+    [("version", 99), ("iterator", -1), ("set_count", -1), ("shuffle_state", None)],
+)
+def test_invalid_reader_progress_is_rejected(tmp_path, field, value):
+    path = _write_system(tmp_path / "data")
+    data = DeepmdData(path)
+    data.get_batch(2)
+    initial = data.state_dict()
+    invalid = copy.deepcopy(initial)
+    invalid[field] = value
+    with pytest.raises(ValueError):
+        data.load_state_dict(invalid)
+    assert data.state_dict() == initial
diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py
index ff11aff68d..cc9b092fad 100644
--- a/source/tests/common/test_examples.py
+++ b/source/tests/common/test_examples.py
@@ -53,7 +53,6 @@
     p_examples / "spin" / "se_e2_a" / "input_tf.json",
     p_examples / "spin" / "se_e2_a" / "input_torch.json",
     p_examples / "spin" / "dpa4" / "input.json",
-    p_examples / "spin" / "dpa4" / "input-deepspin.json",
     p_examples / "spin" / "dpa4c" / "input.json",
     p_examples / "dprc" / "normal" / "input.json",
     p_examples / "dprc" / "pairwise" / "input.json",
@@ -78,6 +77,8 @@
     p_examples / "hessian" / "single_task" / "input.json",
     p_examples / "water" / "se_e2_a" / "input_torch_num_epoch.json",
     p_examples / "lmdb_downsample_data" / "input_lmdb.json",
+    p_examples / "density" / "dpa2" / "input.json",
+    p_examples / "density" / "dpa3" / "input.json",
 )
 
 input_files_multi = (
diff --git a/source/tests/common/test_finetune_utils.py b/source/tests/common/test_finetune_utils.py
index 33547b34d0..622040c2d1 100644
--- a/source/tests/common/test_finetune_utils.py
+++ b/source/tests/common/test_finetune_utils.py
@@ -4,6 +4,8 @@
     deepcopy,
 )
 
+import pytest
+
 from deepmd.utils import (
     finetune,
 )
@@ -156,6 +158,28 @@ def test_finetune_rule_builder_updates_single_task_config_preserving_trainable()
     assert rule.get_has_new_type()
 
 
+def test_finetune_rule_builder_carries_long_range_preserving_trainable():
+    long_range = {"type": "latent_charge", "kernel": {"type": "les", "alpha": 0.7}}
+    pretrained = _model_config(["O", "H"])
+    pretrained["long_range"] = deepcopy(long_range)
+    target = _model_config(["O", "H"], trainable=False)
+    target["long_range"] = {
+        "type": "latent_charge",
+        "charge_net": {"trainable": False},
+        "kernel": {"type": "les", "alpha": 2.0, "trainable": False},
+    }
+
+    updated, _ = finetune.FinetuneRuleBuilder(
+        pretrained,
+        target,
+        change_model_params=True,
+    ).build()
+
+    assert updated["long_range"]["kernel"]["alpha"] == 0.7
+    assert updated["long_range"]["kernel"]["trainable"] is False
+    assert updated["long_range"]["charge_net"]["trainable"] is False
+
+
 def test_finetune_rule_builder_random_fitting_keeps_target_fitting_net():
     pretrained = _model_config(["O", "H"], descriptor_sel=[8, 16], fitting_neuron=[32])
     target = _model_config(["O", "H"], descriptor_sel=[1, 1], fitting_neuron=[2])
@@ -298,22 +322,62 @@ def test_finetune_rule_builder_accepts_compatible_native_spin_transfers():
         ),
     )
     for pretrained, target in compatible_pairs:
-        finetune.FinetuneRuleBuilder(
+        _, rules = finetune.FinetuneRuleBuilder(
             pretrained,
             target,
             change_model_params=False,
         ).build()
+        assert not rules["Default"].convert_to_native_spin
+
+
+def test_finetune_rule_builder_accepts_spin_free_pretraining() -> None:
+    pretrained = _model_config(["Fe", "C"])
+    target = _native_spin_model_config(["Fe", "C"], ["Fe"])
+    target["preset_out_bias"] = {"energy": {"Fe": -1.0, "C": -2.0}}
+    target["bridging_method"] = "zbl"
+    original = deepcopy(target)
+    for change_model_params in (False, True):
+        config, rules = finetune.FinetuneRuleBuilder(
+            pretrained, target, change_model_params=change_model_params
+        ).build()
+        assert rules["Default"].convert_to_native_spin
+        assert config["spin"] == original["spin"]
+        assert config["preset_out_bias"] == original["preset_out_bias"]
+        assert config["bridging_method"] == "zbl"
+        assert target == original
+
+
+def test_spin_free_multitask_branch_requires_finetune_statistics() -> None:
+    pretrained = {"model_dict": {"magnetic": _model_config(["Fe", "C"])}}
+    target = {
+        "model_dict": {"magnetic": _native_spin_model_config(["Fe", "C"], ["Fe"])}
+    }
+    _, rules = finetune.FinetuneRuleBuilder(pretrained, target).build()
+    assert rules["magnetic"].convert_to_native_spin
+    assert not rules["magnetic"].get_resuming()
+
+
+@pytest.mark.parametrize("type_map", [["Fe", "C", "Ni"], ["Fe"], ["C", "Fe"]])
+def test_spin_free_conversion_rejects_changed_type_map(type_map: list[str]) -> None:
+    with pytest.raises(ValueError, match="identical ordered type maps"):
+        finetune.FinetuneRuleBuilder(
+            _model_config(["Fe", "C"]),
+            _native_spin_model_config(type_map, ["Fe"]),
+            change_model_params=False,
+        ).build()
 
 
-def test_finetune_rule_builder_rejects_non_native_spin_pretraining():
+def test_finetune_rule_builder_rejects_virtual_spin_pretraining():
+    pretrained = _model_config(["Fe", "C"])
+    pretrained["spin"] = {"use_spin": [True, False], "scheme": "deepspin"}
     try:
         finetune.FinetuneRuleBuilder(
-            _model_config(["Fe", "C"]),
+            pretrained,
             _native_spin_model_config(["Fe", "C"], [True, False]),
             change_model_params=False,
         ).build()
     except ValueError as exc:
-        assert "requires a native-spin pretrained model" in str(exc)
+        assert "requires a spin-free or native-spin pretrained model" in str(exc)
     else:
         raise AssertionError("expected ValueError")
 
diff --git a/source/tests/common/test_json_schema.py b/source/tests/common/test_json_schema.py
index 843af94da8..eecc27866f 100644
--- a/source/tests/common/test_json_schema.py
+++ b/source/tests/common/test_json_schema.py
@@ -49,6 +49,27 @@ def _assert_valid(validator: Draft202012Validator, data: dict[str, Any]) -> None
     assert not errors, [(list(error.path), error.message) for error in errors]
 
 
+@pytest.mark.parametrize("multi_task", [False, True])
+def test_training_group_schema(
+    validators: dict[bool, Draft202012Validator], multi_task: bool
+) -> None:
+    model = {"preset": "dpa4-nano-v20260911"}
+    if multi_task:
+        model = {"model_dict": {"a": model}}
+    data = _input(model, multi_task)
+    task = data["training"]["data_dict"]["a"] if multi_task else data["training"]
+    task["training_data"] = [
+        {"systems": "first.lmdb", "prob": 0.6},
+        {"systems": ["second.lmdb", "third.lmdb"], "prob": 0.4},
+    ]
+    _assert_valid(validators[multi_task], data)
+
+    task["training_data"][1]["batch_size"] = False
+    assert not validators[multi_task].is_valid(data)
+    task["training_data"] = []
+    assert not validators[multi_task].is_valid(data)
+
+
 def test_preset_user_example(validators: dict[bool, Draft202012Validator]) -> None:
     data = _input(
         {
diff --git a/source/tests/common/test_loggers_training.py b/source/tests/common/test_loggers_training.py
index 97a1874cfe..4620774e3a 100644
--- a/source/tests/common/test_loggers_training.py
+++ b/source/tests/common/test_loggers_training.py
@@ -257,7 +257,8 @@ def test_node_summary_does_not_take_checkpoint_ownership(
 
     output = capsys.readouterr().err
     assert "Computed num_steps=20" in output
-    assert "defaulting to the number of systems per task" in output
+    assert "defaulting to the native data size per task" in output
+    assert "NumPy/HDF5 systems or LMDB frames" in output
     assert task_schedule.model_prob.tolist() == [1 / 3, 2 / 3]
     assert "Resolved checkpoint retention to 2" in output
     assert store.max_keep == 2
diff --git a/source/tests/common/test_path.py b/source/tests/common/test_path.py
index ec3f175c0a..c5fbac3b3e 100644
--- a/source/tests/common/test_path.py
+++ b/source/tests/common/test_path.py
@@ -8,6 +8,9 @@
 import h5py
 import numpy as np
 
+from deepmd.utils.data_system import (
+    process_systems,
+)
 from deepmd.utils.path import (
     DPH5Path,
     DPPath,
@@ -74,3 +77,60 @@ def test_write_operations_are_rejected(self) -> None:
             (self.path / "value").save_numpy(np.ones(1))
         with self.assertRaisesRegex(ValueError, "read-only"):
             (self.path / "group").mkdir()
+
+
+class TestSystemDiscovery(unittest.TestCase):
+    def test_lmdb_paths_and_parent_directories(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            direct = root / "direct.lmdb"
+            direct.mkdir()
+            parent = root / "collection"
+            nested = parent / "nested" / "frames"
+            nested.mkdir(parents=True)
+            (nested / "data.mdb").touch()
+            npy = parent / "npy"
+            npy.mkdir()
+            (npy / "type.raw").touch()
+
+            found = process_systems([str(direct), str(parent)])
+
+            self.assertEqual(found[0], str(direct))
+            self.assertCountEqual(found, [str(direct), str(nested), str(npy)])
+
+    def test_lmdb_pattern_expansion(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            first = root / "train.lmdb"
+            first.mkdir()
+            second = root / "test.lmdb"
+            second.mkdir()
+
+            self.assertEqual(
+                process_systems(str(root), patterns=["train*.lmdb"]),
+                [str(first)],
+            )
+
+    def test_lmdb_is_terminal_even_when_patterns_exclude_it(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            internal = root / "frames.lmdb" / "nested"
+            internal.mkdir(parents=True)
+            (internal / "type.raw").touch()
+            self.assertEqual(process_systems(str(root), patterns=["nested"]), [])
+
+    def test_hdf5_paths_expand_inside_containers(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            container = Path(directory) / "systems.h5"
+            with h5py.File(container, "w") as handle:
+                handle.create_dataset("first/type.raw", data=[0])
+                handle.create_dataset("nested/second/type.raw", data=[0, 1])
+            expected = [f"{container}#/first", f"{container}#/nested/second"]
+            self.assertEqual(process_systems(str(container)), expected)
+            self.assertEqual(
+                process_systems([expected[1], expected[0]]), expected[::-1]
+            )
+            self.assertEqual(
+                process_systems("systems.h5#/nested", root_dir=directory),
+                [expected[1]],
+            )
diff --git a/source/tests/common/test_preset_out_bias.py b/source/tests/common/test_preset_out_bias.py
index 02cadb51bc..8e38405b50 100644
--- a/source/tests/common/test_preset_out_bias.py
+++ b/source/tests/common/test_preset_out_bias.py
@@ -18,11 +18,47 @@
     normalize_preset_out_bias,
     override_assigned_bias,
     preset_out_bias_rows,
+    recover_nonfinite_preset_bias,
     remap_preset_out_bias,
     resolve_preset_out_bias_tables,
 )
 
 
+class TestRecoverNonfinitePresetBias(unittest.TestCase):
+    def test_recovers_only_assigned_nonfinite_cells(self) -> None:
+        stored = np.array(
+            [
+                [[np.nan, 0, 0], [7, 0, 0], [np.inf, 0, 0]],
+                [[1, np.nan, 3], [4, np.inf, 6], [7, 8, -np.inf]],
+                [[np.nan, 0, 0], [1, 0, 0], [2, 0, 0]],
+            ]
+        )
+        original = stored.copy()
+        preset = {
+            "energy": [[-10], [-13.6], [-3]],
+            "tensor": [[10, 11, 12], None, [30, 31, 32]],
+        }
+        rows = recover_nonfinite_preset_bias(
+            preset, ["O", "H", "B"], stored, ["energy", "tensor", "other"], [1, 3, 1]
+        )
+        np.testing.assert_array_equal(rows["energy"], [[-10], [7], [-3]])
+        np.testing.assert_array_equal(
+            rows["tensor"], [[1, 11, 3], [4, np.inf, 6], [7, 8, 32]]
+        )
+        self.assertNotIn("other", rows)
+        np.testing.assert_array_equal(stored, original)
+
+    def test_finite_adjusted_bias_needs_no_recovery(self) -> None:
+        rows = recover_nonfinite_preset_bias(
+            {"energy": [[-10], [-13.6]]},
+            ["O", "H"],
+            np.array([[[7.0], [8.0]]]),
+            ["energy"],
+            [1],
+        )
+        self.assertEqual(rows, {})
+
+
 class TestNormalizePresetOutBias(unittest.TestCase):
     def setUp(self) -> None:
         self.type_map = ["O", "H", "B"]
@@ -109,6 +145,43 @@ def test_errors(self) -> None:
 
 
 class TestBundledTables(unittest.TestCase):
+    def test_domain_cluster(self) -> None:
+        table = load_preset_out_bias_table("domain_cluster")
+        self.assertEqual(
+            set(table), {"Al", "Si", "Ni", "Cu", "Ru", "Pd", "Ag", "Pt", "Au"}
+        )
+        out = normalize_preset_out_bias(
+            {"energy": "domain_cluster"}, ["Au", "Al", "Cu", "H"]
+        )
+        self.assertEqual(
+            out["energy"],
+            [[-901.9377033597224], [-52.96613964130058], [-1306.1758472225818], None],
+        )
+        resolved = resolve_preset_out_bias_tables(
+            {"preset_out_bias": {"energy": "domain_cluster"}}
+        )
+        self.assertEqual(resolved["preset_out_bias"], {"energy": table})
+
+    def test_openmd(self) -> None:
+        table = load_preset_out_bias_table("openmd")
+        self.assertEqual(len(table), 83)
+        out = normalize_preset_out_bias(
+            {"energy": "openmd"}, ["O", "H", "La", "Bi", "Pu"]
+        )
+        self.assertEqual(
+            out["energy"],
+            [
+                [-443.0343602483389],
+                [-12.560490553805243],
+                [-950.8377156904378],
+                [-6345.199411321936],
+                None,
+            ],
+        )
+        config = {"type_map": ["H"], "preset_out_bias": {"energy": "openmd"}}
+        resolved = resolve_preset_out_bias_tables(config)
+        self.assertEqual(resolved["preset_out_bias"], {"energy": table})
+
     def test_name_before_path(self) -> None:
         tables = bundled_preset_out_bias_tables()
         self.assertEqual(load_preset_out_bias_table("omat24"), tables["omat24"])
diff --git a/source/tests/common/test_pytorch_requirements.py b/source/tests/common/test_pytorch_requirements.py
new file mode 100644
index 0000000000..f04b277320
--- /dev/null
+++ b/source/tests/common/test_pytorch_requirements.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""PyTorch extras keep Vesin bindings within their supported versions."""
+
+from pathlib import (
+    Path,
+)
+
+import pytest
+from packaging.requirements import (
+    Requirement,
+)
+
+
+@pytest.mark.parametrize("version", ["2.1.0", "2.2.0", "2.3.0", ""])
+def test_vesin_requirements(version: str, monkeypatch: pytest.MonkeyPatch) -> None:
+    # The PEP 517 backend belongs to the source checkout, not the installed wheel.
+    monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3]))
+    from backend.find_pytorch import (
+        get_pt_requirement,
+    )
+
+    monkeypatch.delenv("PYTORCH_VERSION", raising=False)
+    monkeypatch.setenv("CIBUILDWHEEL", "0")
+    get_pt_requirement.cache_clear()
+    requirements = {
+        requirement.name: requirement
+        for requirement in map(Requirement, get_pt_requirement(version)["torch"])
+    }
+    get_pt_requirement.cache_clear()
+    if version in ("2.1.0", "2.2.0"):
+        assert "vesin" not in requirements
+        assert "vesin-torch" not in requirements
+        return
+
+    assert requirements["vesin"].extras == {"torch"}
+    for name in ("vesin", "vesin-torch"):
+        requirement = requirements[name]
+        assert "0.5.9" not in requirement.specifier
+        assert "0.6" in requirement.specifier
+        for system, machine, supported in (
+            ("Darwin", "x86_64", False),
+            ("Darwin", "arm64", True),
+            ("Linux", "x86_64", True),
+            ("Windows", "AMD64", True),
+        ):
+            assert (
+                requirement.marker.evaluate(
+                    {"platform_system": system, "platform_machine": machine}
+                )
+                is supported
+            )
diff --git a/source/tests/consistent/test_activation.py b/source/tests/consistent/test_activation.py
index 9704ebe60f..8f23a291eb 100644
--- a/source/tests/consistent/test_activation.py
+++ b/source/tests/consistent/test_activation.py
@@ -114,32 +114,6 @@ def test_tf2_consistent_with_ref(self) -> None:
         test = get_activation_fn_dp(self.activation)(input)
         np.testing.assert_allclose(self.ref, to_numpy_array(test), atol=1e-7)
 
-    def test_the_tf2_namespace_name_still_resolves(self) -> None:
-        """Anchor the literal the tf2 dispatch compares against.
-
-        ``xp_erf`` and its neighbours in ``deepmd/dpmodel/array_api.py`` select
-        the TensorFlow path by comparing ``xp.__name__`` against the import
-        path of the vendored namespace -- the idiom this tree uses in
-        ``utils/nlist.py`` and ``utils/default_neighbor_list.py`` too, nine
-        occurrences across three modules.
-
-        A literal like that fails silently if the module is ever renamed or
-        moved: the comparison simply stops matching, the code falls back to the
-        NumPy round-trip, and the gradient is lost again with every
-        forward-value test still passing. This asserts the name resolves to a
-        real module, so a rename breaks a test instead of a derivative.
-
-        Deliberately not gated on ``INSTALLED_TF2``: ``find_spec`` does not
-        import the module, so this runs everywhere, including the ordinary runs
-        where the tf2 cases skip. A guard that skips alongside the thing it
-        guards would protect nothing.
-        """
-        self.assertIsNotNone(
-            find_spec("deepmd._vendors.ndtensorflow"),
-            "the namespace name that array_api.py dispatches on no longer "
-            "resolves; the TensorFlow branches there are now dead code",
-        )
-
     @unittest.skipUnless(INSTALLED_TF2, "TensorFlow 2 is not installed")
     def test_tf2_gradient_consistent_with_ref(self) -> None:
         """The derivative has to survive the dispatch, not only the value.
@@ -262,3 +236,33 @@ def test_pt_expt_consistent_with_ref(self) -> None:
             )
             test = _torch_activation(x, self.activation).detach().cpu().numpy()
             np.testing.assert_allclose(self.ref, test, atol=1e-10)
+
+
+class TestTensorFlowNamespace(unittest.TestCase):
+    """Literal-based TensorFlow dispatch remains reachable without TF installed."""
+
+    def test_the_tf2_namespace_name_still_resolves(self) -> None:
+        """Anchor the literal the tf2 dispatch compares against.
+
+        ``xp_erf`` and its neighbours in ``deepmd/dpmodel/array_api.py`` select
+        the TensorFlow path by comparing ``xp.__name__`` against the import
+        path of the vendored namespace -- the idiom this tree uses in
+        ``utils/nlist.py`` and ``utils/default_neighbor_list.py`` too, nine
+        occurrences across three modules.
+
+        A literal like that fails silently if the module is ever renamed or
+        moved: the comparison simply stops matching, the code falls back to the
+        NumPy round-trip, and the gradient is lost again with every
+        forward-value test still passing. This asserts the name resolves to a
+        real module, so a rename breaks a test instead of a derivative.
+
+        Deliberately not gated on ``INSTALLED_TF2``: ``find_spec`` does not
+        import the module, so this runs everywhere, including the ordinary runs
+        where the tf2 cases skip. A guard that skips alongside the thing it
+        guards would protect nothing.
+        """
+        self.assertIsNotNone(
+            find_spec("deepmd._vendors.ndtensorflow"),
+            "the namespace name that array_api.py dispatches on no longer "
+            "resolves; the TensorFlow branches there are now dead code",
+        )
diff --git a/source/tests/infer/gen_dpa4_spin.py b/source/tests/infer/gen_dpa4_spin.py
index 9cbe25c701..67143150bc 100644
--- a/source/tests/infer/gen_dpa4_spin.py
+++ b/source/tests/infer/gen_dpa4_spin.py
@@ -75,8 +75,7 @@
 # attention + spin-embedding paths, small enough to keep the AOTInductor
 # compile bounded. ``use_spin=[True, False]``: type 0 ("Ni") carries a
 # magnetic moment, type 1 ("O") does not. ``scheme="native"`` selects the
-# NeighborGraph-only spin route (no virtual atoms, unlike the deepspin
-# ``spin_ener`` scheme used by ``gen_spin.py``).
+# NeighborGraph spin route.
 NATIVE_SPIN_CONFIG = {
     "type_map": ["Ni", "O"],
     "descriptor": {
diff --git a/source/tests/infer/gen_dpa4c_canonical.py b/source/tests/infer/gen_dpa4c_canonical.py
new file mode 100644
index 0000000000..0351fd94d1
--- /dev/null
+++ b/source/tests/infer/gen_dpa4c_canonical.py
@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Generate the compressed DPA4C .pt2 test models of the compact canonical lower.
+
+A compressed DPA4C frozen with the compact canonical lower is served by two
+entry points: the host ones, which build the compact canonical graph from
+coordinates, and the device-resident ``compute_canonical_graph_gpu``, which
+LAMMPS + Kokkos calls with the graph and the outputs already on the model
+device. Two archives reach it: an energy model, which ``DeepPotPTExpt``
+serves, and a native-spin model, which ``NativeSpinPTExpt`` serves. The C++
+cases built on them (``dpa4c_canonical_test_common.h``) evaluate the graph
+of the host entry points through the device-resident one, including the
+replay of repeated evaluations of a graph without halo nodes and a charge
+state applied at run time, and check the host entry points against the
+references written here.
+
+Both descriptors are charge-conditioned, so each archive carries the rebuild
+of its frozen tables that ``set_charge_spin`` runs. The compressed operators
+are fp32 CUDA kernels, so the fixtures are generated only where the suite
+has a GPU.
+"""
+
+import copy
+import json
+import os
+import sys
+import zipfile
+
+import numpy as np
+
+# The gen scripts run standalone, outside pytest's package machinery, so the
+# relative imports the test suite uses do not apply here.
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from gen_common import (
+    ensure_inductor_compiler,
+    load_custom_ops,
+    write_expected_ref,
+)
+
+# Small fp32 DPA4C, the configurations the compressed operators specialize on:
+# eight channels and degree two, the profile of the smallest released grade.
+CANONICAL_CONFIG = {
+    "type_map": ["O", "H"],
+    "descriptor": {
+        "type": "dpa4c",
+        "rcut": 4.0,
+        "channels": 8,
+        "lmax": 2,
+        "basis_type": "bessel",
+        "n_radial": 8,
+        "precision": "float32",
+        "seed": 11,
+        "add_chg_spin_ebd": True,
+        "default_chg_spin": [0.0, 1.0],
+    },
+    "fitting_net": {
+        "type": "ener",
+        "neuron": [16, 16],
+        "resnet_dt": False,
+        "activation_function": "silu",
+        "precision": "float32",
+        "seed": 11,
+    },
+}
+
+# The same descriptor with a native spin branch on its first type, which
+# carries a magnetic moment, while the second carries none.
+SPIN_CONFIG = copy.deepcopy(CANONICAL_CONFIG)
+SPIN_CONFIG["type_map"] = ["Ni", "O"]
+SPIN_CONFIG["spin"] = {"use_spin": [True, False], "scheme": "native"}
+
+# Jittered simple cubic lattice of 27 atoms with a spacing of 3 Å in a 9 Å
+# box: every atom has its six nearest neighbours within the 4 Å cutoff, and
+# the box exceeds twice the cutoff, so the single-rank graph has one edge per
+# neighbouring pair and no halo node. The lattice sits half a spacing inside
+# the box, so that every jittered atom lies in it, as the host neighbour list
+# of the C++ suite expects.
+_SPACING = 3.0
+_CELLS = 3
+_NATOMS = _CELLS**3
+_GRID = np.stack(
+    np.meshgrid(*(np.arange(_CELLS),) * 3, indexing="ij"), axis=-1
+).reshape(_NATOMS, 3)
+_RNG = np.random.default_rng(20261016)
+_COORDS = ((_GRID + 0.5) * _SPACING + _RNG.uniform(-0.2, 0.2, (_NATOMS, 3))).reshape(
+    1, _NATOMS, 3
+)
+_ATYPES = (np.arange(_NATOMS) % 2).astype(np.int32)
+_CELL = (np.eye(3) * (_SPACING * _CELLS)).reshape(1, 9)
+# The moments of the spin model, on every atom: the gating of the model, not
+# the input, keeps the second type non-magnetic.
+_SPINS = _RNG.uniform(-0.5, 0.5, (1, _NATOMS, 3))
+
+# The state the archive is frozen with, and one a caller applies at run time.
+_DEFAULT_CHG_SPIN = np.array([[0.0, 1.0]])
+_OTHER_CHG_SPIN = np.array([[1.0, 2.0]])
+
+# Seed of the charge-state head of the energy model, whose output projection
+# starts at zero.
+_HEAD_SEED = 5
+
+# Seed and scale of the zero-initialized parameters of the spin model, which
+# include its spin gate: a closed gate would make every magnetic force vanish.
+_SPIN_REDRAW_SEED = 7
+_SPIN_REDRAW_SCALE = 0.3
+
+
+def _compress(model):
+    """Compress the descriptor of ``model`` as a deployment does."""
+    descriptor = model.atomic_model.fused_decomposition()[0].descriptor
+    descriptor.enable_compression(1.4, table_stride_1=0.01)
+    return model
+
+
+def _build_model():
+    """Build, charge-activate and compress the energy model."""
+    import torch
+
+    from deepmd.pt_expt.model.get_model import (
+        get_model,
+    )
+
+    model = get_model(copy.deepcopy(CANONICAL_CONFIG)).eval()
+    descriptor = model.atomic_model.fused_decomposition()[0].descriptor
+    # The projection of the charge-state head is zero initialized, which
+    # would make the state unobservable and every comparison between states
+    # vacuous.
+    head = descriptor.charge_spin_embedding.network.layers[-1]
+    generator = torch.Generator(device=head.w.device).manual_seed(_HEAD_SEED)
+    with torch.no_grad():
+        head.w.copy_(
+            0.5
+            * torch.randn(
+                head.w.shape,
+                dtype=head.w.dtype,
+                device=head.w.device,
+                generator=generator,
+            )
+        )
+    return _compress(model)
+
+
+def _build_spin_model():
+    """Build the native-spin model, open its zero parameters and compress it."""
+    import torch
+
+    from deepmd.pt_expt.model.get_model import (
+        get_model,
+    )
+
+    model = get_model(copy.deepcopy(SPIN_CONFIG)).eval()
+    generator = torch.Generator(device="cpu").manual_seed(_SPIN_REDRAW_SEED)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            if parameter.numel() == 0 or bool((parameter != 0).any()):
+                continue
+            values = _SPIN_REDRAW_SCALE * torch.randn(
+                parameter.shape, generator=generator, dtype=torch.float64
+            )
+            parameter.copy_(values.to(dtype=parameter.dtype, device=parameter.device))
+    return _compress(model)
+
+
+def _check_metadata(pt2_path: str, spin: bool) -> None:
+    """Assert the archive is the compact canonical lower with a charge state."""
+    with zipfile.ZipFile(pt2_path) as zf:
+        md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8"))
+        names = zf.namelist()
+    print("\n// metadata:")  # noqa: T201
+    print(  # noqa: T201
+        json.dumps(
+            {
+                k: md.get(k)
+                for k in (
+                    "type_map",
+                    "lower_input_kind",
+                    "do_atomic_virial",
+                    "has_chg_spin_ebd",
+                    "default_chg_spin",
+                    "spin_scheme",
+                    "has_comm_artifact",
+                )
+            },
+            indent=2,
+        )
+    )
+    assert md["lower_input_kind"] == "dpa4c_canonical"
+    assert md["do_atomic_virial"] is True
+    # The compact canonical lower takes no condition per call; the state is
+    # frozen into its tables, which ``set_charge_spin`` rebuilds. Without the
+    # rebuild a runtime charge state cannot be served, and the C++ cases built
+    # on it would not reach the replay.
+    assert md.get("charge_state_constants"), (
+        "the archive carries no charge-state rebuild"
+    )
+    assert any(name.endswith("charge_state.pt2") for name in names)
+    if spin:
+        # The dispatch serves a native-spin archive without a with-comm
+        # artifact through NativeSpinPTExpt, whose device entry is the one
+        # the spin cases reach.
+        assert md["spin_scheme"] == "native"
+        assert md.get("has_comm_artifact") is False
+        assert "force_mag" in md["output_keys"]
+
+
+def _write_ref(pt2_path: str, ref_path: str, spin: bool) -> None:
+    """Evaluate the archive under both states and write the C++ sidecar."""
+    from deepmd.infer import (
+        DeepPot,
+    )
+
+    dp = DeepPot(pt2_path)
+    assert dp.has_spin is spin
+    # The system itself, which the C++ cases read back rather than restate.
+    system = {
+        "coord": _COORDS[0],
+        "atype": _ATYPES.astype(np.float64),
+        "box": _CELL[0],
+    }
+    if spin:
+        system["spin"] = _SPINS[0]
+    sections = {"system": system}
+    energies = {}
+    for name, charge_spin in (
+        ("default_state", _DEFAULT_CHG_SPIN),
+        ("other_state", _OTHER_CHG_SPIN),
+    ):
+        if spin:
+            e, f, v, ae, av, fm, _ = dp.eval(
+                _COORDS,
+                _CELL,
+                _ATYPES,
+                atomic=True,
+                spin=_SPINS,
+                charge_spin=charge_spin,
+            )
+        else:
+            e, f, v, ae, av = dp.eval(
+                _COORDS, _CELL, _ATYPES, atomic=True, charge_spin=charge_spin
+            )
+        assert np.all(np.isfinite(e)) and np.all(np.isfinite(f))
+        energies[name] = float(e[0, 0])
+        print(f"// {name} total energy: {e[0, 0]:.18e}")  # noqa: T201
+        sections[name] = {
+            "expected_e": ae[0, :, 0],
+            "expected_f": f[0],
+            "expected_tot_v": v[0],
+            "expected_atom_v": av[0],
+        }
+        if spin:
+            magnetic = fm[0].reshape(_NATOMS, 3)
+            # A closed spin branch yields no magnetic force at all, which
+            # would leave the magnetic outputs of the C++ cases unchecked;
+            # the gating of the model zeroes those of the second type.
+            assert np.abs(magnetic[_ATYPES == 0]).max() > 1e-6, (
+                "the magnetic force vanishes on the magnetic atoms"
+            )
+            assert np.abs(magnetic[_ATYPES == 1]).max() == 0.0, (
+                "the magnetic force does not vanish on the non-magnetic atoms"
+            )
+            sections[name]["expected_fm"] = fm[0]
+    # The C++ cases compare evaluations under the two states with each other.
+    assert abs(energies["other_state"] - energies["default_state"]) > 1e-4, (
+        "the two charge states must differ, or a replay that kept the "
+        "constants of the first state would still match"
+    )
+    write_expected_ref(
+        ref_path,
+        sections=sections,
+        source_script="source/tests/infer/gen_dpa4c_canonical.py",
+    )
+    print(f"Wrote {ref_path}")  # noqa: T201
+
+
+def main():
+    from deepmd.pt_expt.utils.serialization import (
+        deserialize_to_file as pt_expt_deserialize_to_file,
+    )
+
+    ensure_inductor_compiler()
+    load_custom_ops()
+
+    base_dir = os.path.dirname(__file__)
+    for stem, build, spin in (
+        ("deeppot_dpa4c_canonical", _build_model, False),
+        ("deeppot_dpa4c_canonical_spin", _build_spin_model, True),
+    ):
+        pt2_path = os.path.join(base_dir, f"{stem}.pt2")
+        print(f"Exporting to {pt2_path} (lower_kind='dpa4c_canonical') ...")  # noqa: T201
+        pt_expt_deserialize_to_file(
+            pt2_path,
+            {"model": build().serialize()},
+            lower_kind="dpa4c_canonical",
+            do_atomic_virial=True,
+        )
+        print("Export done.")  # noqa: T201
+        _check_metadata(pt2_path, spin)
+        _write_ref(pt2_path, os.path.join(base_dir, f"{stem}.expected"), spin)
+    print("\nDone!")  # noqa: T201
+
+
+if __name__ == "__main__":
+    main()
diff --git a/source/tests/infer/gen_dpa4c_long_range.py b/source/tests/infer/gen_dpa4c_long_range.py
new file mode 100644
index 0000000000..3c71034dd5
--- /dev/null
+++ b/source/tests/infer/gen_dpa4c_long_range.py
@@ -0,0 +1,323 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Generate the DPA4C long-range .pt2 test models and reference values.
+
+A latent-charge long-range module sums over every pair of atoms of a complete
+non-periodic system, and over the images of the cell of a periodic one when
+it serves periodic systems, so the C++ interface hands the model what the sum
+needs beyond the neighbor list: the graph lower of an uncompressed model
+receives the pair list, or the mesh of the cell, that the interface builds
+from the positions of the atoms, and the compact canonical lower of a
+compressed model reads the positions and the cell themselves. One archive
+covers each lower and each kind of module:
+
+- ``deeppot_dpa4c_long_range.pt2`` and
+  ``deeppot_dpa4c_long_range_periodic.pt2``: uncompressed, fp64, graph lower,
+  which compiles on the CPU; the periodic one needs torch >= 2.12, whose
+  AOTInductor packages the mesh evaluation;
+- ``deeppot_dpa4c_long_range_compressed.pt2`` and
+  ``deeppot_dpa4c_long_range_periodic_compressed.pt2``: compressed, fp32,
+  compact canonical lower, whose operators are CUDA only; they are generated
+  for the CUDA variant of the C++ suite (``DP_VARIANT=cuda``).
+
+Every model conditions the descriptor on the charge state and constrains the
+latent charges to the total charge that state names, so a state other than
+the default changes the long-range energy as well as the descriptor. The
+references are those of the Python interface on each archive. For a
+non-periodic module they are taken without a cell, at the default state and
+at ``[1, 2]`` (sections ``default`` and ``explicit``); for a periodic module,
+in the cell :data:`BOX` and without a cell, at the neutral default state
+(sections ``periodic`` and ``cluster``). The six atoms form two groups about
+9 Å apart, beyond the cutoff of the descriptor, which only the long-range
+module couples; in the cell, the first group meets its images across the
+boundary.
+
+``deeppot_dpa4c_long_range_periodic_mesh.expected`` holds the mesh topology
+that the Python builder forms for the periodic case, against which the C++
+suite checks the mesh it builds for the graph lower; it needs no archive.
+"""
+
+import copy
+import dataclasses
+import json
+import os
+import sys
+import zipfile
+
+import numpy as np
+
+# The gen scripts run standalone, outside pytest's package machinery, so the
+# relative imports the test suite uses do not apply here.
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from dpa4_fixtures import (
+    jitter_zero_arrays,
+)
+from gen_common import (
+    ensure_inductor_compiler,
+    load_custom_ops,
+    write_expected_ref,
+)
+
+LONG_RANGE = {
+    "type": "latent_charge",
+    "channels": 1,
+    "charge_net": {"neuron": [8, 8], "seed": 9},
+    "kernel": {"type": "les", "alpha": 0.8},
+    "charge_constraint": {"source": "charge_spin"},
+}
+
+# The uncompressed model is fp64, so that its references hold to the
+# precision of the double C++ interface.
+CONFIG = {
+    "type_map": ["O", "H"],
+    "descriptor": {
+        "type": "dpa4c",
+        "rcut": 4.0,
+        "channels": 8,
+        "lmax": 2,
+        "n_radial": 4,
+        "precision": "float64",
+        "seed": 7,
+        "add_chg_spin_ebd": True,
+        "default_chg_spin": [0.0, 1.0],
+    },
+    "fitting_net": {"neuron": [8, 8], "precision": "float64", "seed": 7},
+    "long_range": LONG_RANGE,
+}
+
+# Compression takes an fp32 descriptor and, for the compact canonical lower,
+# a fitting that the fused fitting operator evaluates.
+COMPRESSED_CONFIG = {
+    "type_map": ["O", "H"],
+    "descriptor": {
+        "type": "dpa4c",
+        "rcut": 4.0,
+        "channels": 8,
+        "lmax": 2,
+        "n_radial": 8,
+        "precision": "float32",
+        "seed": 7,
+        "add_chg_spin_ebd": True,
+        "default_chg_spin": [0.0, 1.0],
+    },
+    "fitting_net": {
+        "neuron": [32, 32],
+        "activation_function": "silu",
+        "resnet_dt": False,
+        "precision": "float32",
+        "seed": 7,
+    },
+    "long_range": LONG_RANGE,
+}
+
+# Jitters the zero-initialized arrays, among them the charge-state head of the
+# descriptor, so that the charge state reaches the descriptor as well as the
+# constraint.
+_JITTER_SEED = 20261018
+
+_ATYPES = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32)
+_COORDS = np.array(
+    [
+        [12.83, 2.56, 2.18],
+        [12.09, 2.87, 2.74],
+        [0.25, 3.32, 1.68],
+        [3.36, 3.00, 1.81],
+        [3.51, 2.51, 2.60],
+        [4.27, 3.22, 1.56],
+    ],
+    dtype=np.float64,
+).reshape(1, 6, 3)
+_OTHER_CHG_SPIN = np.array([[1.0, 2.0]], dtype=np.float64)
+# A triclinic cell, row-major with the lattice vectors as rows.
+BOX = np.array(
+    [[13.0, 0.0, 0.0], [0.6, 12.5, 0.0], [-0.4, 0.8, 13.5]], dtype=np.float64
+).reshape(1, 9)
+
+
+def _periodic(config: dict) -> dict:
+    """The configuration with a long-range module that serves periodic systems."""
+    config = copy.deepcopy(config)
+    config["long_range"]["periodic"] = True
+    return config
+
+
+def _model_dict(config: dict) -> dict:
+    """Build the dpmodel of a configuration and jitter it in place of its zeros."""
+    from deepmd.dpmodel.model.model import (
+        get_model,
+    )
+
+    model = get_model(copy.deepcopy(config))
+    return jitter_zero_arrays(model.serialize(), np.random.default_rng(_JITTER_SEED))
+
+
+def _check_metadata(pt2_path: str, lower: str, periodic: bool = False) -> None:
+    """Assert the archive declares its lower and the long-range input it needs."""
+    with zipfile.ZipFile(pt2_path) as zf:
+        md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8"))
+    assert md["lower_input_kind"] == lower, md["lower_input_kind"]
+    expected = ["all_pairs", "pme_mesh"] if periodic else ["all_pairs"]
+    assert md["topologies"] == expected, md["topologies"]
+    if lower == "graph":
+        assert md["dim_chg_spin"] == 2
+    else:
+        # The compressed descriptor folds the charge state; the archive
+        # carries the rebuild of every constant that holds it: the four
+        # charge-state artifacts of the descriptor and the total charge of
+        # the constraint.
+        assert md["dim_chg_spin"] == 0
+        assert len(md["charge_state_constants"]) == 5
+        # The fused long-range operator serves the charge network of one
+        # channel and reads the source CSR for the per-atom virial alone.
+        assert md["canonical_requires_source_csr"] is False
+
+
+def _write_ref(pt2_path: str, ref_path: str, periodic: bool = False) -> None:
+    """Evaluate the reference cases on the archive and write the C++ sidecar."""
+    from deepmd.infer import (
+        DeepPot,
+    )
+
+    dp = DeepPot(pt2_path)
+    if periodic:
+        cases = (("periodic", BOX, None), ("cluster", None, None))
+    else:
+        cases = (("default", None, None), ("explicit", None, _OTHER_CHG_SPIN))
+    sections = {}
+    energies = {}
+    for name, box, charge_spin in cases:
+        condition = {} if charge_spin is None else {"charge_spin": charge_spin}
+        e, f, v, ae, av = dp.eval(_COORDS, box, _ATYPES, atomic=True, **condition)
+        assert np.all(np.isfinite(e)) and np.all(np.isfinite(f))
+        energies[name] = float(e[0, 0])
+        print(f"// {name} total energy: {e[0, 0]:.18e}")  # noqa: T201
+        sections[name] = {
+            "expected_e": ae[0, :, 0],
+            "expected_f": f[0],
+            "expected_v": av[0],
+        }
+    # The C++ cases compare the two cases, which needs a real difference.
+    first, second = (name for name, _, _ in cases)
+    assert abs(energies[second] - energies[first]) > 1e-3, (
+        "the two cases must differ, or a call that ignored the cell or the "
+        "requested state would still match"
+    )
+    write_expected_ref(
+        ref_path,
+        sections=sections,
+        source_script="source/tests/infer/gen_dpa4c_long_range.py",
+    )
+    print(f"Wrote {ref_path}")  # noqa: T201
+
+
+def _write_mesh_ref(config: dict, ref_path: str) -> None:
+    """Write the mesh topology of the reference frame, built in Python, for the C++ twin.
+
+    The graph lower of a periodic archive reads the mesh that the C++
+    interface builds from the cell at the spacing of the archive. The
+    reference is the mesh of the Python builder at that spacing, in the cell
+    :data:`BOX` and without a cell (sections ``periodic`` and ``cluster``).
+    It needs no archive, so it is written by every PyTorch release.
+    """
+    from deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph import (
+        build_mesh_grid,
+        empty_mesh_grid,
+    )
+    from deepmd.pt_expt.model.model import (
+        BaseModel,
+    )
+
+    model = BaseModel.deserialize(_model_dict(config))
+    module = model.atomic_model.long_range
+    module.use_mesh_evaluation()
+    spacing = module.topology_params()["pme_mesh"]["spacing"]
+    coord = _COORDS[0]
+    sections = {}
+    for name, mesh in (
+        ("periodic", build_mesh_grid(BOX, spacing, coord)),
+        ("cluster", empty_mesh_grid(1, coord)),
+    ):
+        sections[name] = {"spacing": np.array([spacing])}
+        for field in dataclasses.fields(mesh):
+            sections[name][field.name] = np.asarray(
+                getattr(mesh, field.name), dtype=np.float64
+            )
+    write_expected_ref(
+        ref_path,
+        sections=sections,
+        source_script="source/tests/infer/gen_dpa4c_long_range.py",
+    )
+    print(f"Wrote {ref_path}")  # noqa: T201
+
+
+def _export_compressed(pt2_path: str, config: dict) -> None:
+    """Compress the fp32 model and export it for the CUDA device."""
+    from deepmd.pt_expt.model.model import (
+        BaseModel,
+    )
+    from deepmd.pt_expt.utils.serialization import (
+        deserialize_to_file,
+    )
+
+    model = BaseModel.deserialize(_model_dict(config))
+    model.atomic_model.enable_compression(0.5)
+    data = {
+        "model": model.serialize(),
+        "model_def_script": config,
+        "backend": "dpmodel",
+        "software": "deepmd-kit",
+        "version": "3.0.0",
+    }
+    print(f"Exporting to {pt2_path} ...")  # noqa: T201
+    deserialize_to_file(pt2_path, data, do_atomic_virial=True, lower_kind="auto")
+
+
+def main() -> None:
+    from deepmd.pt_expt.utils.serialization import (
+        deserialize_to_file,
+        mesh_freeze_supported,
+    )
+
+    ensure_inductor_compiler()
+    load_custom_ops()
+
+    base_dir = os.path.dirname(__file__)
+    for suffix, periodic in (("", False), ("_periodic", True)):
+        config = _periodic(CONFIG) if periodic else CONFIG
+        stem = os.path.join(base_dir, f"deeppot_dpa4c_long_range{suffix}")
+        data = {
+            "model": _model_dict(config),
+            "model_def_script": config,
+            "backend": "dpmodel",
+            "software": "deepmd-kit",
+            "version": "3.0.0",
+        }
+        if periodic:
+            _write_mesh_ref(config, f"{stem}_mesh.expected")
+        if periodic and not mesh_freeze_supported():
+            # The C++ and LAMMPS suites skip the archive this torch cannot
+            # package.
+            print(f"Skipping {stem}.pt2: the mesh evaluation needs torch >= 2.12")  # noqa: T201
+        else:
+            print(f"Exporting to {stem}.pt2 (lower_kind='graph') ...")  # noqa: T201
+            deserialize_to_file(
+                f"{stem}.pt2", data, do_atomic_virial=True, lower_kind="graph"
+            )
+            _check_metadata(f"{stem}.pt2", "graph", periodic)
+            _write_ref(f"{stem}.pt2", f"{stem}.expected", periodic)
+
+        if os.environ.get("DP_VARIANT") == "cuda":
+            compressed = f"{stem}_compressed"
+            compressed_config = (
+                _periodic(COMPRESSED_CONFIG) if periodic else COMPRESSED_CONFIG
+            )
+            _export_compressed(f"{compressed}.pt2", compressed_config)
+            _check_metadata(f"{compressed}.pt2", "dpa4c_canonical", periodic)
+            _write_ref(f"{compressed}.pt2", f"{compressed}.expected", periodic)
+
+    print("\nDone!")  # noqa: T201
+
+
+if __name__ == "__main__":
+    main()
diff --git a/source/tests/infer/gen_dpa4c_spin_lmp.py b/source/tests/infer/gen_dpa4c_spin_lmp.py
new file mode 100644
index 0000000000..bdd13c0599
--- /dev/null
+++ b/source/tests/infer/gen_dpa4c_spin_lmp.py
@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Generate a minimal native-spin DPA4C graph ``.pt2`` for LAMMPS ``dpa4spin``.
+
+Unlike ``gen_dpa4c_spin.py``, this archive carries no fparam, aparam, or
+charge/spin condition so ``pair_style dpa4spin`` can load it without extra
+keywords. Used by ``source/lmp/tests/test_dpa4spin_magnetic_force_convention.py``.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+import copy
+import json
+import os
+import sys
+import zipfile
+
+import numpy as np
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from dpa4_fixtures import (
+    jitter_zero_arrays,
+)
+from gen_common import (
+    ensure_inductor_compiler,
+    load_custom_ops,
+)
+
+CFG = {
+    "type_map": ["Ni", "O"],
+    "descriptor": {
+        "type": "dpa4c",
+        "rcut": 4.0,
+        "channels": 8,
+        "lmax": 2,
+        "n_radial": 4,
+        "precision": "float64",
+        "seed": 7,
+    },
+    "fitting_net": {"neuron": [8, 8], "precision": "float64", "seed": 7},
+    "spin": {"use_spin": [True, False], "scheme": "native"},
+}
+
+_JITTER_SEED = 20260817
+OUT = os.path.join(os.path.dirname(__file__), "deeppot_dpa4c_spin_lmp.pt2")
+
+
+def main() -> None:
+    from torch._inductor import config as inductor_config
+
+    from deepmd.dpmodel.model.model import (
+        get_model,
+    )
+    from deepmd.pt_expt.utils.serialization import (
+        deserialize_to_file,
+    )
+
+    ensure_inductor_compiler()
+    load_custom_ops()
+    model = get_model(copy.deepcopy(CFG))
+    md = jitter_zero_arrays(model.serialize(), np.random.default_rng(_JITTER_SEED))
+    data = {
+        "model": md,
+        "model_def_script": CFG,
+        "backend": "dpmodel",
+        "software": "deepmd-kit",
+        "version": "3.0.0",
+    }
+    print(f"Exporting {OUT}")  # noqa: T201
+    # PyTorch 2.11's single-threaded CPU SIMD gather codegen can use an
+    # index temporary before declaring it. This tiny fixture checks adapter
+    # magnetic-force conventions, so scalar codegen avoids that compiler bug
+    # without changing production exports or CUDA compilation.
+    with inductor_config.patch({"cpp.simdlen": 1}):
+        deserialize_to_file(OUT, data, do_atomic_virial=True, lower_kind="graph")
+    with zipfile.ZipFile(OUT) as zf:
+        meta = json.loads(zf.read("model/extra/metadata.json").decode())
+    print(  # noqa: T201
+        json.dumps(
+            {
+                k: meta.get(k)
+                for k in (
+                    "spin_scheme",
+                    "has_comm_artifact",
+                    "dim_fparam",
+                    "dim_aparam",
+                    "dim_chg_spin",
+                    "lower_input_kind",
+                )
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/source/tests/infer/gen_dpa4c_spin_lmp_kokkos.py b/source/tests/infer/gen_dpa4c_spin_lmp_kokkos.py
new file mode 100644
index 0000000000..80f6475bf6
--- /dev/null
+++ b/source/tests/infer/gen_dpa4c_spin_lmp_kokkos.py
@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Generate the compressed CUDA fixture for LAMMPS ``dpa4spin/kk``.
+
+The Kokkos adapter requires the compact canonical graph ABI, which is only
+available for a compressed float32 DPA4C model with a fused fitting network.
+Keep this archive separate from the float64 CPU graph fixture. Both carry
+native spin and no frame, atomic, or charge/spin conditioning parameters.
+
+Run with a CUDA-enabled DeePMD build and a visible GPU. A CPU-only invocation
+fails rather than silently exporting an archive the Kokkos adapter cannot use.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+import copy
+import json
+import os
+import zipfile
+
+import gen_dpa4c_spin_lmp as host_fixture
+import numpy as np
+from gen_common import (
+    ensure_inductor_compiler,
+    load_custom_ops,
+)
+
+CFG = copy.deepcopy(host_fixture.CFG)
+CFG["descriptor"]["precision"] = "float32"
+CFG["fitting_net"]["precision"] = "float32"
+# The fused fitting kernel does not support per-layer residual timesteps.
+CFG["fitting_net"]["resnet_dt"] = False
+OUT = os.path.join(os.path.dirname(__file__), "deeppot_dpa4c_spin_lmp_kokkos.pt2")
+
+
+def _build_model_data() -> dict:
+    """Build reproducible weights and compression tables on the CPU."""
+    from deepmd.dpmodel.model.model import (
+        get_model,
+    )
+    from deepmd.pt_expt.kernels.dpa4c.graph_compress import (
+        mega_eligible,
+    )
+    from deepmd.pt_expt.kernels.graph_fitting import (
+        fitting_eligible,
+    )
+    from deepmd.pt_expt.model.model import (
+        BaseModel,
+    )
+
+    model_dict = host_fixture.jitter_zero_arrays(
+        get_model(copy.deepcopy(CFG)).serialize(),
+        np.random.default_rng(host_fixture._JITTER_SEED),
+    )
+    model = BaseModel.deserialize(model_dict).to("cpu").eval()
+    descriptor = model.get_descriptor()
+    assert mega_eligible(descriptor), "the fixture needs a compiled DPA4C profile"
+    assert fitting_eligible(model.get_fitting_net()), "the fixture needs fused fitting"
+    # Jitter before compression so the tabulated snapshot retains the live
+    # native-spin branch, including its derivative at exactly zero spin.
+    descriptor.enable_compression(min_nbor_dist=0.5)
+    assert descriptor.compress
+    return {
+        "model": model.serialize(),
+        "model_def_script": copy.deepcopy(CFG),
+        "backend": "pt_expt",
+        "software": "deepmd-kit",
+        "version": "3.0.0",
+    }
+
+
+def _check_metadata(path: str) -> None:
+    """Reject an archive that would exercise a different deployment route."""
+    with zipfile.ZipFile(path) as zf:
+        metadata = json.loads(zf.read("model/extra/metadata.json"))
+    expected = {
+        "type_map": ["Ni", "O"],
+        "spin_scheme": "native",
+        "use_spin": [True, False],
+        "has_comm_artifact": False,
+        "dim_fparam": 0,
+        "dim_aparam": 0,
+        "dim_chg_spin": 0,
+        "lower_input_kind": "dpa4c_canonical",
+        "canonical_index_dtype": "auto",
+        "canonical_index_storage": "byte_words",
+        "graph_edge_dtype": "float32",
+    }
+    for key, value in expected.items():
+        assert metadata.get(key) == value, (key, metadata.get(key), value)
+    for key in ("energy", "force", "force_mag", "virial", "atom_virial"):
+        assert key in metadata["output_keys"], key
+    print(json.dumps({key: metadata[key] for key in expected}, indent=2))  # noqa: T201
+
+
+def main() -> None:
+    from deepmd.pt_expt.kernels.dpa4c.canonical import (
+        canonical_model_eligible,
+        op_available,
+    )
+    from deepmd.pt_expt.model.model import (
+        BaseModel,
+    )
+    from deepmd.pt_expt.utils.env import (
+        DEVICE,
+    )
+    from deepmd.pt_expt.utils.serialization import (
+        deserialize_to_file,
+    )
+
+    if DEVICE.type != "cuda":
+        raise RuntimeError(
+            "The dpa4spin/kk fixture requires a CUDA-enabled DeePMD build "
+            "and a visible GPU (without DEVICE=cpu)."
+        )
+    ensure_inductor_compiler()
+    load_custom_ops()
+    if not op_available():
+        raise RuntimeError(
+            "The compressed canonical DPA4C CUDA operators are required."
+        )
+    data = _build_model_data()
+    assert canonical_model_eligible(BaseModel.deserialize(data["model"]))
+    print(f"Exporting {OUT}")  # noqa: T201
+    deserialize_to_file(OUT, data, do_atomic_virial=True, lower_kind="dpa4c_canonical")
+    _check_metadata(OUT)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/source/tests/pd/model/test_atomic_model_global_stat.py b/source/tests/pd/model/test_atomic_model_global_stat.py
index bacec7e682..023bafb277 100644
--- a/source/tests/pd/model/test_atomic_model_global_stat.py
+++ b/source/tests/pd/model/test_atomic_model_global_stat.py
@@ -4,6 +4,9 @@
 from pathlib import (
     Path,
 )
+from unittest.mock import (
+    patch,
+)
 
 import h5py
 import numpy as np
@@ -360,6 +363,42 @@ def cvt_ret(x):
         out_bias, _ = md0._fetch_out_stat(["foo", "bar"])
         np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias)
         np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias)
+
+        # Standalone change-bias explicitly ignores the checkpoint preset.
+        adjusted_bias = foo_bias + 7.0
+        md0._store_out_stat({"foo": to_paddle_tensor(adjusted_bias)})
+        forward = md0._get_forward_wrapper_func()
+        seen_biases = []
+
+        def check_forward(*args, **kwargs):
+            bias, _ = md0._fetch_out_stat(["foo"])
+            seen_biases.append(to_numpy_array(bias["foo"]).copy())
+            return forward(*args, **kwargs)
+
+        with patch.object(md0, "_get_forward_wrapper_func", return_value=check_forward):
+            BaseAtomicModel.change_out_bias(
+                md0,
+                self.merged_output_stat,
+                bias_adjust_mode="change-by-statistic",
+                preset_out_bias={},
+            )
+        self.assertTrue(seen_biases)
+        for seen_bias in seen_biases:
+            np.testing.assert_almost_equal(seen_bias, adjusted_bias)
+        out_bias, _ = md0._fetch_out_stat(["foo", "bar"])
+        np.testing.assert_almost_equal(
+            to_numpy_array(out_bias["foo"]), np.array([2.0, -5.0]).reshape(2, 1)
+        )
+        assert not np.allclose(to_numpy_array(out_bias["bar"]), bar_bias)
+
+        # The override does not remove the model's configured preset.
+        np.testing.assert_almost_equal(md0.preset_out_bias["foo"], foo_bias)
+        BaseAtomicModel.change_out_bias(
+            md0, self.merged_output_stat, bias_adjust_mode="change-by-statistic"
+        )
+        out_bias, _ = md0._fetch_out_stat(["foo", "bar"])
+        np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias)
+        np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias)
         # 4. a preset that leaves an observed type unassigned is rejected
         md1 = DPAtomicModel(
             ds,
diff --git a/source/tests/pd/model/test_autodiff.py b/source/tests/pd/model/test_autodiff.py
index 8442844a24..78060cb481 100644
--- a/source/tests/pd/model/test_autodiff.py
+++ b/source/tests/pd/model/test_autodiff.py
@@ -27,10 +27,7 @@
 from .test_permutation import (
     model_dpa1,
     model_dpa2,
-    model_hybrid,
     model_se_e2_a,
-    model_spin,
-    model_zbl,
 )
 
 
@@ -209,52 +206,3 @@ def setUp(self) -> None:
         model_params = copy.deepcopy(model_dpa2)
         self.type_split = True
         self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelDPAUniVirial(unittest.TestCase, VirialTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_dpa2)
-        self.type_split = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelHybridForce(unittest.TestCase, ForceTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_hybrid)
-        self.type_split = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelHybridVirial(unittest.TestCase, VirialTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_hybrid)
-        self.type_split = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelZBLForce(unittest.TestCase, ForceTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_zbl)
-        self.type_split = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelZBLVirial(unittest.TestCase, VirialTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_zbl)
-        self.type_split = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelSpinSeAForce(unittest.TestCase, ForceTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_spin)
-        self.type_split = False
-        self.test_spin = True
-        self.model = get_model(model_params).to(env.DEVICE)
diff --git a/source/tests/pd/model/test_dpa3.py b/source/tests/pd/model/test_dpa3.py
index 7b3c8ba18f..a42069b0c3 100644
--- a/source/tests/pd/model/test_dpa3.py
+++ b/source/tests/pd/model/test_dpa3.py
@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import itertools
 import unittest
+from unittest.mock import (
+    patch,
+)
 
 import numpy as np
 import paddle
@@ -36,6 +39,125 @@ class TestDescrptDPA3(unittest.TestCase, TestCaseSingleFrameWithNlist):
     def setUp(self) -> None:
         TestCaseSingleFrameWithNlist.setUp(self)
 
+    def test_non_trainable_parameters_stop_gradient(self) -> None:
+        """A non-trainable descriptor must disable Paddle autograd parameters."""
+        # Paddle 3.0 does not expose the newer ``requires_grad`` compatibility
+        # alias. Ignore writes to that alias so this test retains minimum-version
+        # semantics when the CI environment uses a newer Paddle release.
+        requires_grad_alias = property(
+            lambda parameter: not parameter.stop_gradient,
+            lambda _parameter, _value: None,
+        )
+        with patch.object(
+            paddle.Tensor, "requires_grad", requires_grad_alias, create=True
+        ):
+            descriptor = DescrptDPA3(
+                self.nt,
+                repflow=RepFlowArgs(
+                    n_dim=4,
+                    e_dim=4,
+                    a_dim=4,
+                    nlayers=1,
+                    e_sel=2,
+                    a_sel=1,
+                    axis_neuron=2,
+                ),
+                trainable=False,
+                add_chg_spin_ebd=True,
+                seed=GLOBAL_SEED,
+            )
+            serialized_descriptor = descriptor.serialize()
+            deserialized_descriptor = DescrptDPA3.deserialize(serialized_descriptor)
+
+        # Optional charge/spin layers must serialize their effective frozen
+        # state so cross-backend deserializers do not re-enable gradients.
+        self.assertFalse(serialized_descriptor["chg_embedding"]["trainable"])
+        self.assertFalse(serialized_descriptor["spin_embedding"]["trainable"])
+        self.assertFalse(serialized_descriptor["mix_cs_mlp"]["trainable"])
+
+        for stage, checked_descriptor in (
+            ("constructed", descriptor),
+            ("deserialized", deserialized_descriptor),
+        ):
+            with self.subTest(stage=stage):
+                parameters = list(checked_descriptor.named_parameters())
+                self.assertTrue(parameters)
+                parameters_with_grad = [
+                    name
+                    for name, parameter in parameters
+                    if not parameter.stop_gradient
+                ]
+                self.assertEqual([], parameters_with_grad)
+
+    def test_non_trainable_change_type_map_stays_frozen(self) -> None:
+        """Type remapping must not make a frozen embedding trainable again."""
+        descriptor = DescrptDPA3(
+            self.nt,
+            repflow=RepFlowArgs(
+                n_dim=4,
+                e_dim=4,
+                a_dim=4,
+                nlayers=1,
+                e_sel=2,
+                a_sel=1,
+                axis_neuron=2,
+            ),
+            trainable=False,
+            type_map=["O", "H"],
+            seed=GLOBAL_SEED,
+        )
+
+        descriptor.change_type_map(["H", "O"])
+
+        parameters_with_grad = [
+            name
+            for name, parameter in descriptor.named_parameters()
+            if not parameter.stop_gradient
+        ]
+        self.assertEqual([], parameters_with_grad)
+
+    def test_share_params_rejects_mismatched_trainable(self) -> None:
+        """Sharing must not let one descriptor rewrite another's gradients."""
+        repflow = RepFlowArgs(
+            n_dim=4,
+            e_dim=4,
+            a_dim=4,
+            nlayers=1,
+            e_sel=2,
+            a_sel=1,
+            axis_neuron=2,
+        )
+        for shared_level in (0, 1):
+            with self.subTest(shared_level=shared_level):
+                trainable_descriptor = DescrptDPA3(
+                    self.nt,
+                    repflow=repflow,
+                    trainable=True,
+                    seed=GLOBAL_SEED,
+                )
+                frozen_descriptor = DescrptDPA3(
+                    self.nt,
+                    repflow=repflow,
+                    trainable=False,
+                    seed=GLOBAL_SEED,
+                )
+
+                with self.assertRaisesRegex(ValueError, "same trainable setting"):
+                    trainable_descriptor.share_params(
+                        frozen_descriptor, shared_level=shared_level
+                    )
+
+                # Validation happens before either the embedding or repflow
+                # layers can be aliased to the other descriptor.
+                self.assertIsNot(
+                    trainable_descriptor.type_embedding,
+                    frozen_descriptor.type_embedding,
+                )
+                self.assertIsNot(
+                    trainable_descriptor.repflows,
+                    frozen_descriptor.repflows,
+                )
+
     def test_consistency(
         self,
     ) -> None:
diff --git a/source/tests/pd/model/test_force_grad.py b/source/tests/pd/model/test_force_grad.py
deleted file mode 100644
index d1387771f1..0000000000
--- a/source/tests/pd/model/test_force_grad.py
+++ /dev/null
@@ -1,108 +0,0 @@
-# SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
-import json
-import unittest
-from pathlib import (
-    Path,
-)
-
-import numpy as np
-import paddle
-
-from deepmd.pd.model.model import (
-    get_model,
-)
-from deepmd.pd.utils import (
-    env,
-)
-from deepmd.utils.data import (
-    DeepmdData,
-)
-
-from ...seed import (
-    GLOBAL_SEED,
-)
-
-
-class CheckSymmetry(DeepmdData):
-    def __init__(
-        self,
-        sys_path: str,
-        type_map: list[str] | None = None,
-    ):
-        super().__init__(sys_path=sys_path, type_map=type_map)
-        self.add("energy", 1, atomic=False, must=False, high_prec=True)
-        self.add("force", 3, atomic=True, must=False, high_prec=False)
-        self.add("virial", 9, atomic=False, must=False, high_prec=False)
-
-    def get_disturb(self, index, atom_index, axis_index, delta):
-        for i in range(
-            len(self.dirs) + 1
-        ):  # note: if different sets can be merged, prefix sum is unused to calculate
-            if index < self.prefix_sum[i]:
-                break
-        frames = self._load_set(self.dirs[i - 1])
-        tmp = copy.deepcopy(frames["coord"].reshape(self.nframes, -1, 3))
-        tmp[:, atom_index, axis_index] += delta
-        frames["coord"] = tmp
-        frame = self._get_subdata(frames, index - self.prefix_sum[i - 1])
-        frame = self.reformat_data_torch(frame)
-        return frame
-
-
-def get_data(batch):
-    inputs = {}
-    for key in ["coord", "atype", "box"]:
-        inputs[key] = batch[key].unsqueeze(0).to(env.DEVICE)
-    return inputs
-
-
-class TestForceGrad(unittest.TestCase):
-    def setUp(self):
-        with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin:
-            self.config = json.load(fin)
-        data_file = [str(Path(__file__).parent / "water/data/data_0")]
-        self.config["training"]["training_data"]["systems"] = data_file
-        self.config["training"]["validation_data"]["systems"] = data_file
-        self.system_index = 0
-        self.batch_index = 0
-        self.get_dataset(self.system_index, self.batch_index)
-        self.get_model()
-
-    def get_model(self):
-        self.model = get_model(self.config["model"]).to(env.DEVICE)
-
-    def get_dataset(self, system_index=0, batch_index=0):
-        systems = self.config["training"]["training_data"]["systems"]
-        rcut = self.config["model"]["descriptor"]["rcut"]
-        sel = self.config["model"]["descriptor"]["sel"]
-        sec = paddle.cumsum(paddle.to_tensor(sel), axis=0)
-        type_map = self.config["model"]["type_map"]
-        self.dpdatasystem = CheckSymmetry(
-            sys_path=systems[system_index], type_map=type_map
-        )
-        self.origin_batch = self.dpdatasystem.get_item_paddle(batch_index)
-
-    @unittest.skip("it can be replaced by autodiff")
-    def test_force_grad(self, threshold=1e-2, delta0=1e-6, seed=20):
-        rng = np.random.default_rng(GLOBAL_SEED)
-        result0 = self.model(**get_data(self.origin_batch))
-        np.random.default_rng(seed)
-        errors = np.zeros((self.dpdatasystem.natoms, 3))
-        for atom_index in range(self.dpdatasystem.natoms):
-            for axis_index in range(3):
-                delta = rng.random() * delta0
-                disturb_batch = self.dpdatasystem.get_disturb(
-                    self.batch_index, atom_index, axis_index, delta
-                )
-                disturb_result = self.model(**get_data(disturb_batch))
-                disturb_force = -(disturb_result["energy"] - result0["energy"]) / delta
-                disturb_error = (
-                    result0["force"][0, atom_index, axis_index] - disturb_force
-                )
-                errors[atom_index, axis_index] = disturb_error.detach().cpu().numpy()
-        self.assertTrue(np.abs(errors).max() < threshold, msg=str(np.abs(errors).max()))
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/source/tests/pd/model/test_forward_lower.py b/source/tests/pd/model/test_forward_lower.py
index 1d924e2d3d..a60e10acf5 100644
--- a/source/tests/pd/model/test_forward_lower.py
+++ b/source/tests/pd/model/test_forward_lower.py
@@ -26,8 +26,6 @@
     model_dpa1,
     model_dpa2,
     model_se_e2_a,
-    model_spin,
-    model_zbl,
 )
 
 dtype = paddle.float64
@@ -160,47 +158,5 @@ def setUp(self) -> None:
         self.model = get_model(model_params).to(env.DEVICE)
 
 
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelZBL(unittest.TestCase, ForwardLowerTest):
-    def setUp(self) -> None:
-        self.prec = 1e-10
-        model_params = copy.deepcopy(model_zbl)
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelSpinSeA(unittest.TestCase, ForwardLowerTest):
-    def setUp(self) -> None:
-        self.prec = 1e-10
-        model_params = copy.deepcopy(model_spin)
-        self.test_spin = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelSpinDPA1(unittest.TestCase, ForwardLowerTest):
-    def setUp(self) -> None:
-        self.prec = 1e-10
-        model_params = copy.deepcopy(model_spin)
-        model_params["descriptor"] = copy.deepcopy(model_dpa1)["descriptor"]
-        # double sel for virtual atoms to avoid large error
-        model_params["descriptor"]["sel"] *= 2
-        self.test_spin = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelSpinDPA2(unittest.TestCase, ForwardLowerTest):
-    def setUp(self) -> None:
-        self.prec = 1e-10
-        model_params = copy.deepcopy(model_spin)
-        model_params["descriptor"] = copy.deepcopy(model_dpa2)["descriptor"]
-        # double sel for virtual atoms to avoid large error
-        model_params["descriptor"]["repinit"]["nsel"] *= 2
-        model_params["descriptor"]["repformer"]["nsel"] *= 2
-        self.test_spin = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pd/model/test_model.py b/source/tests/pd/model/test_model.py
index f45e02391b..6afece5ef4 100644
--- a/source/tests/pd/model/test_model.py
+++ b/source/tests/pd/model/test_model.py
@@ -411,12 +411,15 @@ def test_consistency(self) -> None:
             .numpy(),
         )
         # self.assertIsNone(model_predict_1.get("atom_virial", None))
+        # Atomic virials can be close to zero after cancellation; allow
+        # float64 roundoff without relaxing the default relative tolerance.
         np.testing.assert_allclose(
             head_dict["atom_virial"],
             p_atomic_virial.reshape(head_dict["atom_virial"].shape)
             .cpu()
             .detach()
             .numpy(),
+            atol=1e-14,
         )
         optimizer = paddle.optimizer.Adam(cur_lr, parameters=my_model.parameters())
         optimizer.clear_grad()
diff --git a/source/tests/pd/model/test_permutation.py b/source/tests/pd/model/test_permutation.py
index 048cb2ed4b..1b414b50ef 100644
--- a/source/tests/pd/model/test_permutation.py
+++ b/source/tests/pd/model/test_permutation.py
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import copy
-import os
 import unittest
 
 import numpy as np
@@ -20,8 +19,6 @@
     eval_model,
 )
 
-CUR_DIR = os.path.dirname(__file__)
-
 dtype = paddle.float64
 
 model_se_e2_a = {
@@ -44,87 +41,6 @@
     "data_stat_nbatch": 20,
 }
 
-model_dos = {
-    "type_map": ["O", "H", "B"],
-    "descriptor": {
-        "type": "se_e2_a",
-        "sel": [46, 92, 4],
-        "rcut_smth": 0.50,
-        "rcut": 4.00,
-        "neuron": [25, 50, 100],
-        "resnet_dt": False,
-        "axis_neuron": 16,
-        "seed": 1,
-    },
-    "fitting_net": {
-        "neuron": [24, 24, 24],
-        "resnet_dt": True,
-        "seed": 1,
-        "type": "dos",
-        "numb_dos": 250,
-    },
-    "data_stat_nbatch": 20,
-}
-
-model_zbl = {
-    "type_map": ["O", "H", "B"],
-    "use_srtab": f"{CUR_DIR}/water/data/zbl_tab_potential/H2O_tab_potential.txt",
-    "smin_alpha": 0.1,
-    "sw_rmin": 0.2,
-    "sw_rmax": 4.0,
-    "descriptor": {
-        "type": "se_atten",
-        "sel": 40,
-        "rcut_smth": 0.5,
-        "rcut": 4.0,
-        "neuron": [25, 50, 100],
-        "axis_neuron": 16,
-        "attn": 64,
-        "attn_layer": 2,
-        "attn_dotr": True,
-        "attn_mask": False,
-        "activation_function": "tanh",
-        "scaling_factor": 1.0,
-        "normalize": False,
-        "temperature": 1.0,
-        "set_davg_zero": True,
-        "type_one_side": True,
-        "seed": 1,
-    },
-    "fitting_net": {
-        "neuron": [24, 24, 24],
-        "resnet_dt": True,
-        "seed": 1,
-    },
-    "data_stat_nbatch": 20,
-}
-
-
-model_spin = {
-    "type_map": ["O", "H", "B"],
-    "descriptor": {
-        "type": "se_e2_a",
-        "sel": [46, 92, 4],
-        "rcut_smth": 0.50,
-        "rcut": 4.00,
-        "neuron": [25, 50, 100],
-        "resnet_dt": False,
-        "axis_neuron": 16,
-        "seed": 1,
-    },
-    "fitting_net": {
-        "neuron": [24, 24, 24],
-        "resnet_dt": True,
-        "seed": 1,
-    },
-    "data_stat_nbatch": 20,
-    "spin": {
-        "use_spin": [True, False, False],
-        "virtual_scale": [0.3140],
-        "_comment": " that's all",
-    },
-}
-
 model_dpa2 = {
     "type_map": ["O", "H", "B"],
     "descriptor": {
@@ -289,73 +205,6 @@
 }
 
 
-model_hybrid = {
-    "type_map": ["O", "H", "B"],
-    "descriptor": {
-        "type": "hybrid",
-        "list": [
-            {
-                "type": "se_atten",
-                "sel": 120,
-                "rcut_smth": 0.5,
-                "rcut": 6.0,
-                "neuron": [25, 50, 100],
-                "axis_neuron": 16,
-                "attn": 128,
-                "attn_layer": 0,
-                "attn_dotr": True,
-                "attn_mask": False,
-                "activation_function": "tanh",
-                "scaling_factor": 1.0,
-                "normalize": True,
-                "temperature": 1.0,
-                "seed": 1,
-            },
-            {
-                "type": "dpa2",
-                "repinit": {
-                    "rcut": 6.0,
-                    "rcut_smth": 2.0,
-                    "nsel": 30,
-                    "neuron": [2, 4, 8],
-                    "axis_neuron": 4,
-                    "activation_function": "tanh",
-                },
-                "repformer": {
-                    "rcut": 4.0,
-                    "rcut_smth": 0.5,
-                    "nsel": 10,
-                    "nlayers": 12,
-                    "g1_dim": 8,
-                    "g2_dim": 5,
-                    "attn2_hidden": 3,
-                    "attn2_nhead": 1,
-                    "attn1_hidden": 5,
-                    "attn1_nhead": 1,
-                    "axis_neuron": 4,
-                    "update_h2": False,
-                    "update_g1_has_conv": True,
-                    "update_g1_has_grrg": True,
-                    "update_g1_has_drrd": True,
-                    "update_g1_has_attn": True,
-                    "update_g2_has_g1g1": True,
-                    "update_g2_has_attn": True,
-                    "attn2_has_gate": True,
-                },
-                "seed": 1,
-                "add_tebd_to_repinit_out": False,
-            },
-        ],
-    },
-    "fitting_net": {
-        "neuron": [240, 240, 240],
-        "resnet_dt": True,
-        "seed": 1,
-        "_comment": " that's all",
-    },
-    "_comment": " that's all",
-}
-
 model_property = {
     "type_map": ["H", "C", "N", "O"],
     "descriptor": {
@@ -440,14 +289,6 @@ def setUp(self) -> None:
         self.model = get_model(model_params).to(env.DEVICE)
 
 
-@unittest.skip("Skip for not implemented yet")
-class TestDOSModelSeA(unittest.TestCase, PermutationTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_dos)
-        self.type_split = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
 class TestEnergyModelDPA1(unittest.TestCase, PermutationTest):
     def setUp(self) -> None:
         model_params = copy.deepcopy(model_dpa1)
@@ -462,65 +303,5 @@ def setUp(self) -> None:
         self.model = get_model(model_params).to(env.DEVICE)
 
 
-@unittest.skip("Skip for not implemented yet")
-class TestForceModelDPA2(unittest.TestCase, PermutationTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_dpa2)
-        model_params["fitting_net"]["type"] = "direct_force_ener"
-        self.type_split = True
-        self.test_virial = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelHybrid(unittest.TestCase, PermutationTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_hybrid)
-        self.type_split = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestForceModelHybrid(unittest.TestCase, PermutationTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_hybrid)
-        model_params["fitting_net"]["type"] = "direct_force_ener"
-        self.type_split = True
-        self.test_virial = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelZBL(unittest.TestCase, PermutationTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_zbl)
-        self.type_split = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelSpinSeA(unittest.TestCase, PermutationTest):
-    def setUp(self) -> None:
-        model_params = copy.deepcopy(model_spin)
-        self.type_split = False
-        self.test_spin = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-# class TestEnergyFoo(unittest.TestCase):
-#   def test(self):
-#     model_params = model_dpau
-#     self.model = EnergyModelDPAUni(model_params).to(env.DEVICE)
-
-#     natoms = 5
-#     cell = paddle.rand([3, 3], dtype=dtype)
-#     cell = (cell + cell.T) + 5. * paddle.eye(3)
-#     coord = paddle.rand([natoms, 3], dtype=dtype)
-#     coord = paddle.matmul(coord, cell)
-#     atype = paddle.to_tensor([0, 0, 0, 1, 1])
-#     idx_perm = [1, 0, 4, 3, 2]
-#     ret0 = infer_model(self.model, coord, cell, atype, type_split=True)
-
-
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pd/model/test_rot.py b/source/tests/pd/model/test_rot.py
index 84a0d3d724..71f786cb24 100644
--- a/source/tests/pd/model/test_rot.py
+++ b/source/tests/pd/model/test_rot.py
@@ -19,13 +19,9 @@
     eval_model,
 )
 from .test_permutation import (  # model_dpau,
-    model_dos,
     model_dpa1,
     model_dpa2,
-    model_hybrid,
     model_se_e2_a,
-    model_spin,
-    model_zbl,
 )
 
 dtype = paddle.float64
@@ -161,14 +157,6 @@ def setUp(self):
         self.model = get_model(model_params).to(env.DEVICE)
 
 
-@unittest.skip("Skip for not implemented yet")
-class TestDOSModelSeA(unittest.TestCase, RotTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_dos)
-        self.type_split = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
 class TestEnergyModelDPA1(unittest.TestCase, RotTest):
     def setUp(self):
         model_params = copy.deepcopy(model_dpa1)
@@ -183,50 +171,5 @@ def setUp(self):
         self.model = get_model(model_params).to(env.DEVICE)
 
 
-@unittest.skip("Skip for not implemented yet")
-class TestForceModelDPA2(unittest.TestCase, RotTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_dpa2)
-        model_params["fitting_net"]["type"] = "direct_force_ener"
-        self.type_split = True
-        self.test_virial = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelHybrid(unittest.TestCase, RotTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_hybrid)
-        self.type_split = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestForceModelHybrid(unittest.TestCase, RotTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_hybrid)
-        model_params["fitting_net"]["type"] = "direct_force_ener"
-        self.type_split = True
-        self.test_virial = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelZBL(unittest.TestCase, RotTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_zbl)
-        self.type_split = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelSpinSeA(unittest.TestCase, RotTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_spin)
-        self.type_split = False
-        self.test_spin = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pd/model/test_trans.py b/source/tests/pd/model/test_trans.py
index f050596996..6fde3d6687 100644
--- a/source/tests/pd/model/test_trans.py
+++ b/source/tests/pd/model/test_trans.py
@@ -19,13 +19,9 @@
     eval_model,
 )
 from .test_permutation import (  # model_dpau,
-    model_dos,
     model_dpa1,
     model_dpa2,
-    model_hybrid,
     model_se_e2_a,
-    model_spin,
-    model_zbl,
 )
 
 dtype = paddle.float64
@@ -95,14 +91,6 @@ def setUp(self):
         self.model = get_model(model_params).to(env.DEVICE)
 
 
-@unittest.skip("Skip for not implemented yet")
-class TestDOSModelSeA(unittest.TestCase, TransTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_dos)
-        self.type_split = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
 class TestEnergyModelDPA1(unittest.TestCase, TransTest):
     def setUp(self):
         model_params = copy.deepcopy(model_dpa1)
@@ -117,50 +105,5 @@ def setUp(self):
         self.model = get_model(model_params).to(env.DEVICE)
 
 
-@unittest.skip("Skip for not implemented yet")
-class TestForceModelDPA2(unittest.TestCase, TransTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_dpa2)
-        model_params["fitting_net"]["type"] = "direct_force_ener"
-        self.type_split = True
-        self.test_virial = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelHybrid(unittest.TestCase, TransTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_hybrid)
-        self.type_split = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestForceModelHybrid(unittest.TestCase, TransTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_hybrid)
-        model_params["fitting_net"]["type"] = "direct_force_ener"
-        self.type_split = True
-        self.test_virial = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelZBL(unittest.TestCase, TransTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_zbl)
-        self.type_split = False
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyModelSpinSeA(unittest.TestCase, TransTest):
-    def setUp(self):
-        model_params = copy.deepcopy(model_spin)
-        self.type_split = False
-        self.test_spin = True
-        self.model = get_model(model_params).to(env.DEVICE)
-
-
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pd/model/test_unused_params.py b/source/tests/pd/model/test_unused_params.py
deleted file mode 100644
index bf92171da1..0000000000
--- a/source/tests/pd/model/test_unused_params.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
-import unittest
-
-import paddle
-
-from deepmd.pd.model.model import (
-    get_model,
-)
-from deepmd.pd.utils import (
-    env,
-)
-
-from ...seed import (
-    GLOBAL_SEED,
-)
-from ..common import (
-    eval_model,
-)
-from .test_permutation import (
-    model_dpa2,
-)
-
-dtype = paddle.float64
-
-
-@unittest.skip("paddle do not support unpacking grad_fn.next_functions")
-class TestUnusedParamsDPA2(unittest.TestCase):
-    def test_unused(self):
-        import itertools
-
-        for conv, drrd, grrg, attn1, g1g1, attn2, h2 in itertools.product(
-            [True],
-            [True],
-            [True],
-            [True],
-            [True],
-            [True],
-            [True],
-        ):
-            if (not drrd) and (not grrg) and h2:
-                # skip the case h2 is not envolved
-                continue
-            if (not grrg) and (not conv):
-                # skip the case g2 is not envolved
-                continue
-            model = copy.deepcopy(model_dpa2)
-            model["descriptor"]["repformer"]["nlayers"] = 2
-            # model["descriptor"]["combine_grrg"] = cmbg2
-            model["descriptor"]["repformer"]["update_g1_has_conv"] = conv
-            model["descriptor"]["repformer"]["update_g1_has_drrd"] = drrd
-            model["descriptor"]["repformer"]["update_g1_has_grrg"] = grrg
-            model["descriptor"]["repformer"]["update_g1_has_attn"] = attn1
-            model["descriptor"]["repformer"]["update_g2_has_g1g1"] = g1g1
-            model["descriptor"]["repformer"]["update_g2_has_attn"] = attn2
-            model["descriptor"]["repformer"]["update_h2"] = h2
-            model["fitting_net"]["neuron"] = [12, 12, 12]
-            self._test_unused(model)
-
-    def _test_unused(self, model_params):
-        self.model = get_model(model_params).to(env.DEVICE)
-        natoms = 5
-        generator = paddle.seed(GLOBAL_SEED)
-        cell = paddle.rand([3, 3], dtype=dtype).to(device=env.DEVICE)
-        cell = (cell + cell.T) + 5.0 * paddle.eye(3).to(device=env.DEVICE)
-        coord = paddle.rand([natoms, 3], dtype=dtype).to(device=env.DEVICE)
-        coord = paddle.matmul(coord, cell)
-        atype = paddle.to_tensor([0, 0, 0, 1, 1]).to(env.DEVICE)
-        idx_perm = [1, 0, 4, 3, 2]
-        result_0 = eval_model(self.model, coord.unsqueeze(0), cell.unsqueeze(0), atype)
-        test_keys = ["energy", "force", "virial"]
-        ret0 = {key: result_0[key].squeeze(0) for key in test_keys}
-
-        # use computation graph to find all contributing tensors
-        def get_contributing_params(y, top_level=True):
-            nf = y.grad_fn.next_functions if top_level else y.next_functions
-            for f, _ in nf:
-                try:
-                    yield f.variable
-                except AttributeError:
-                    pass  # node has no tensor
-                if f is not None:
-                    yield from get_contributing_params(f, top_level=False)
-
-        contributing_parameters = set(get_contributing_params(ret0["energy"]))
-        all_parameters = set(self.model.parameters())
-        non_contributing = all_parameters - contributing_parameters
-        self.assertEqual(len(non_contributing), 0)
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/source/tests/pd/test_change_bias.py b/source/tests/pd/test_change_bias.py
index 2d87b739ff..36c2f0290c 100644
--- a/source/tests/pd/test_change_bias.py
+++ b/source/tests/pd/test_change_bias.py
@@ -140,6 +140,29 @@ def test_change_bias_with_user_defined(self):
         )
         np.testing.assert_allclose(updated_bias.numpy(), expected_bias.numpy())
 
+    def test_change_bias_ignores_checkpoint_preset(self):
+        checkpoint = paddle.load(str(self.model_path))
+        preset = {"energy": [1.0e6] * len(self.config["model"]["type_map"])}
+        checkpoint["model"]["_extra_state"]["model_params"]["preset_out_bias"] = preset
+        paddle.save(checkpoint, str(self.model_path))
+        run_dp(
+            f"dp --pd change-bias {self.model_path!s} -s {self.data_file[0]} -o {self.model_path_data_bias!s}"
+        )
+        updated = paddle.load(str(self.model_path_data_bias))
+        model_params = updated["model"]["_extra_state"]["model_params"]
+        self.assertEqual(model_params["preset_out_bias"], preset)
+        wrapper = ModelWrapper(get_model_for_wrapper(model_params))
+        wrapper.set_state_dict(updated["model"])
+        expected_model = model_change_out_bias(
+            self.trainer.wrapper.model["Default"],
+            self.sampled,
+            _bias_adjust_mode="change-by-statistic",
+        )
+        np.testing.assert_allclose(
+            wrapper.model["Default"].get_out_bias().numpy(),
+            expected_model.get_out_bias().numpy(),
+        )
+
     def tearDown(self):
         for f in os.listdir("."):
             if f.startswith("change-bias-model") and f.endswith(".pd"):
diff --git a/source/tests/pd/test_dp_show.py b/source/tests/pd/test_dp_show.py
deleted file mode 100644
index 9ce3834e55..0000000000
--- a/source/tests/pd/test_dp_show.py
+++ /dev/null
@@ -1,221 +0,0 @@
-# SPDX-License-Identifier: LGPL-3.0-or-later
-import io
-import json
-import os
-import shutil
-import unittest
-from contextlib import (
-    redirect_stderr,
-)
-from copy import (
-    deepcopy,
-)
-from pathlib import (
-    Path,
-)
-
-from deepmd.pd.entrypoints.main import (
-    get_trainer,
-)
-from deepmd.pd.utils.multi_task import (
-    preprocess_shared_params,
-)
-
-from .common import (
-    run_dp,
-)
-from .model.test_permutation import (
-    model_se_e2_a,
-)
-
-
-@unittest.skip("Skip until solving cuda error 709 in jit.save")
-class TestSingleTaskModel(unittest.TestCase):
-    def setUp(self):
-        input_json = str(Path(__file__).parent / "water/se_atten.json")
-        with open(input_json) as f:
-            self.config = json.load(f)
-        os.environ["FLAGS_prim_enable_dynamic"] = "1"
-        os.environ["FLAGS_enable_pir_api"] = "1"
-        self.config["training"]["numb_steps"] = 1
-        self.config["training"]["save_freq"] = 1
-        data_file = [str(Path(__file__).parent / "water/data/single")]
-        self.config["training"]["training_data"]["systems"] = data_file
-        self.config["training"]["validation_data"]["systems"] = data_file
-        self.config["model"] = deepcopy(model_se_e2_a)
-        self.config["model"]["type_map"] = ["O", "H", "Au"]
-        trainer = get_trainer(deepcopy(self.config))
-        trainer.run()
-        run_dp("dp --pd freeze")
-
-    def test_checkpoint(self):
-        INPUT = "model.pd"
-        ATTRIBUTES = "type-map descriptor fitting-net"
-        with redirect_stderr(io.StringIO()) as f:
-            run_dp(f"dp --pd show {INPUT} {ATTRIBUTES}")
-        results = f.getvalue().split("\n")[:-1]
-        assert "This is a singletask model" in results[-4]
-        assert "The type_map is ['O', 'H', 'Au']" in results[-3]
-        assert (
-            "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0"
-        ) in results[-2]
-        assert (
-            "The fitting_net parameter is {'neuron': [24, 24, 24], 'resnet_dt': True, 'seed': 1}"
-            in results[-1]
-        )
-
-    @unittest.skip(
-        "Paddle do not support dp --pd show frozen models(.json and .pdiparams file), "
-        "will be supported in the future."
-    )
-    def test_frozen_model(self):
-        INPUT = "frozen_model.json"
-        ATTRIBUTES = "type-map descriptor fitting-net"
-        with redirect_stderr(io.StringIO()) as f:
-            run_dp(f"dp --pd show {INPUT} {ATTRIBUTES}")
-        results = f.getvalue().split("\n")[:-1]
-        assert "This is a singletask model" in results[-4]
-        assert "The type_map is ['O', 'H', 'Au']" in results[-3]
-        assert (
-            "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0"
-        ) in results[-2]
-        assert (
-            "The fitting_net parameter is {'neuron': [24, 24, 24], 'resnet_dt': True, 'seed': 1}"
-            in results[-1]
-        )
-
-    def test_checkpoint_error(self):
-        INPUT = "model.pd"
-        ATTRIBUTES = "model-branch type-map descriptor fitting-net"
-        with self.assertRaisesRegex(
-            RuntimeError, "The 'model-branch' option requires a multitask model"
-        ):
-            run_dp(f"dp --pd show {INPUT} {ATTRIBUTES}")
-
-    def tearDown(self):
-        for f in os.listdir("."):
-            if f.startswith("model") and f.endswith("pd"):
-                os.remove(f)
-            if f in ["lcurve.out", "frozen_model.pd", "output.txt", "checkpoint"]:
-                os.remove(f)
-            if f in ["stat_files"]:
-                shutil.rmtree(f)
-
-
-@unittest.skip("Skip until solving cuda error 709 in jit.save")
-class TestMultiTaskModel(unittest.TestCase):
-    def setUp(self):
-        input_json = str(Path(__file__).parent / "water/multitask.json")
-        with open(input_json) as f:
-            self.config = json.load(f)
-        self.config["model"]["shared_dict"]["my_descriptor"] = model_se_e2_a[
-            "descriptor"
-        ]
-        data_file = [str(Path(__file__).parent / "water/data/data_0")]
-        self.stat_files = "se_e2_a"
-        os.makedirs(self.stat_files, exist_ok=True)
-        self.config["training"]["data_dict"]["model_1"]["training_data"]["systems"] = (
-            data_file
-        )
-        self.config["training"]["data_dict"]["model_1"]["validation_data"][
-            "systems"
-        ] = data_file
-        self.config["training"]["data_dict"]["model_1"]["stat_file"] = (
-            f"{self.stat_files}/model_1"
-        )
-        self.config["training"]["data_dict"]["model_2"]["training_data"]["systems"] = (
-            data_file
-        )
-        self.config["training"]["data_dict"]["model_2"]["validation_data"][
-            "systems"
-        ] = data_file
-        self.config["training"]["data_dict"]["model_2"]["stat_file"] = (
-            f"{self.stat_files}/model_2"
-        )
-        self.config["model"]["model_dict"]["model_1"]["fitting_net"] = {
-            "neuron": [1, 2, 3],
-            "seed": 678,
-        }
-        self.config["model"]["model_dict"]["model_2"]["fitting_net"] = {
-            "neuron": [9, 8, 7],
-            "seed": 1111,
-        }
-        self.config["training"]["numb_steps"] = 1
-        self.config["training"]["save_freq"] = 1
-        self.origin_config = deepcopy(self.config)
-        self.config["model"], self.shared_links = preprocess_shared_params(
-            self.config["model"]
-        )
-        trainer = get_trainer(deepcopy(self.config), shared_links=self.shared_links)
-        trainer.run()
-        run_dp("dp --pd freeze --model-branch model_1")
-
-    def test_checkpoint(self):
-        INPUT = "model.ckpt.pd"
-        ATTRIBUTES = "model-branch type-map descriptor fitting-net"
-        with redirect_stderr(io.StringIO()) as f:
-            run_dp(f"dp --pd show {INPUT} {ATTRIBUTES}")
-        results = f.getvalue().split("\n")[:-1]
-        assert "This is a multitask model" in results[-8]
-        assert (
-            "Available model branches are ['model_1', 'model_2', 'RANDOM'], "
-            "where 'RANDOM' means using a randomly initialized fitting net."
-            in results[-7]
-        )
-        assert "The type_map of branch model_1 is ['O', 'H', 'B']" in results[-6]
-        assert "The type_map of branch model_2 is ['O', 'H', 'B']" in results[-5]
-        assert (
-            "model_1"
-            and "'type': 'se_e2_a'"
-            and "'sel': [46, 92, 4]"
-            and "'rcut_smth': 0.5"
-        ) in results[-4]
-        assert (
-            "model_2"
-            and "'type': 'se_e2_a'"
-            and "'sel': [46, 92, 4]"
-            and "'rcut_smth': 0.5"
-        ) in results[-3]
-        assert (
-            "The fitting_net parameter of branch model_1 is {'neuron': [1, 2, 3], 'seed': 678}"
-            in results[-2]
-        )
-        assert (
-            "The fitting_net parameter of branch model_2 is {'neuron': [9, 8, 7], 'seed': 1111}"
-            in results[-1]
-        )
-
-    @unittest.skip(
-        "Paddle do not support dp --pd show frozen models(.json and .pdiparams file), "
-        "will be supported in the future."
-    )
-    def test_frozen_model(self):
-        INPUT = "frozen_model.json"
-        ATTRIBUTES = "type-map descriptor fitting-net"
-        with redirect_stderr(io.StringIO()) as f:
-            run_dp(f"dp --pd show {INPUT} {ATTRIBUTES}")
-        results = f.getvalue().split("\n")[:-1]
-        assert "This is a singletask model" in results[-4]
-        assert "The type_map is ['O', 'H', 'B']" in results[-3]
-        assert (
-            "'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut_smth': 0.5"
-        ) in results[-2]
-        assert (
-            "The fitting_net parameter is {'neuron': [1, 2, 3], 'seed': 678}"
-            in results[-1]
-        )
-
-    def tearDown(self):
-        for f in os.listdir("."):
-            if f.startswith("model") and f.endswith("pd"):
-                os.remove(f)
-            if f in [
-                "lcurve.out",
-                "frozen_model.json",
-                "frozen_model.pdiparams",
-                "checkpoint",
-                "output.txt",
-            ]:
-                os.remove(f)
-            if f in ["stat_files", self.stat_files]:
-                shutil.rmtree(f)
diff --git a/source/tests/pd/test_finetune.py b/source/tests/pd/test_finetune.py
index 769ea6f6d3..ad04674b60 100644
--- a/source/tests/pd/test_finetune.py
+++ b/source/tests/pd/test_finetune.py
@@ -44,11 +44,8 @@
 )
 
 from .model.test_permutation import (
-    model_dos,
     model_dpa1,
     model_dpa2,
-    model_se_e2_a,
-    model_zbl,
 )
 
 energy_data_requirement = [
@@ -294,53 +291,6 @@ def tearDown(self) -> None:
                 shutil.rmtree(f)
 
 
-class TestEnergyModelSeA(FinetuneTest, unittest.TestCase):
-    def setUp(self) -> None:
-        input_json = str(Path(__file__).parent / "water/se_atten.json")
-        with open(input_json) as f:
-            self.config = json.load(f)
-        self.data_file = [str(Path(__file__).parent / "water/data/single")]
-        self.config["training"]["training_data"]["systems"] = self.data_file
-        self.config["training"]["validation_data"]["systems"] = self.data_file
-        self.config["model"] = deepcopy(model_se_e2_a)
-        self.config["training"]["numb_steps"] = 1
-        self.config["training"]["save_freq"] = 1
-        self.mixed_types = False
-        self.testkey = None
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyZBLModelSeA(FinetuneTest, unittest.TestCase):
-    def setUp(self) -> None:
-        input_json = str(Path(__file__).parent / "water/se_atten.json")
-        with open(input_json) as f:
-            self.config = json.load(f)
-        self.data_file = [str(Path(__file__).parent / "water/data/single")]
-        self.config["training"]["training_data"]["systems"] = self.data_file
-        self.config["training"]["validation_data"]["systems"] = self.data_file
-        self.config["model"] = deepcopy(model_zbl)
-        self.config["training"]["numb_steps"] = 1
-        self.config["training"]["save_freq"] = 1
-        self.mixed_types = False
-        self.testkey = None
-
-
-@unittest.skip("Skip for not implemented yet")
-class TestEnergyDOSModelSeA(FinetuneTest, unittest.TestCase):
-    def setUp(self) -> None:
-        input_json = str(Path(__file__).parent / "dos/input.json")
-        with open(input_json) as f:
-            self.config = json.load(f)
-        self.data_file = [str(Path(__file__).parent / "dos/data/global_system")]
-        self.config["training"]["training_data"]["systems"] = self.data_file
-        self.config["training"]["validation_data"]["systems"] = self.data_file
-        self.config["model"] = deepcopy(model_dos)
-        self.config["training"]["numb_steps"] = 1
-        self.config["training"]["save_freq"] = 1
-        self.mixed_types = False
-        self.testkey = "dos"
-
-
 class TestEnergyModelDPA1(FinetuneTest, unittest.TestCase):
     def setUp(self) -> None:
         input_json = str(Path(__file__).parent / "water/se_atten.json")
diff --git a/source/tests/pd/test_lr.py b/source/tests/pd/test_lr.py
deleted file mode 100644
index deef3e2630..0000000000
--- a/source/tests/pd/test_lr.py
+++ /dev/null
@@ -1,176 +0,0 @@
-# SPDX-License-Identifier: LGPL-3.0-or-later
-import unittest
-
-import numpy as np
-import tensorflow.compat.v1 as tf
-
-tf.disable_eager_execution()
-
-from deepmd.dpmodel.utils.learning_rate import (
-    LearningRateExp,
-)
-from deepmd.pd.utils.learning_rate import (
-    LearningRateWSD,
-)
-from deepmd.tf.utils.learning_rate import (
-    LearningRateSchedule,
-)
-
-
-class TestLearningRate(unittest.TestCase):
-    def setUp(self):
-        self.start_lr = 0.001
-        self.stop_lr = 3.51e-8
-        # decay_steps will be auto-adjusted if >= stop_steps
-        self.decay_steps = np.arange(400, 501, 100)
-        self.stop_steps = np.arange(500, 1600, 500)
-
-    def test_consistency(self):
-        for decay_step in self.decay_steps:
-            for stop_step in self.stop_steps:
-                self.decay_step = decay_step
-                self.stop_step = stop_step
-                self.judge_it()
-                self.decay_rate_pd()
-
-    def judge_it(self):
-        base_lr = LearningRateSchedule(
-            {
-                "type": "exp",
-                "start_lr": self.start_lr,
-                "stop_lr": self.stop_lr,
-                "decay_steps": self.decay_step,
-            }
-        )
-        g = tf.Graph()
-        with g.as_default():
-            global_step = tf.placeholder(shape=[], dtype=tf.int32)
-            t_lr = base_lr.build(global_step, self.stop_step)
-
-        my_lr = LearningRateExp(
-            start_lr=self.start_lr,
-            stop_lr=self.stop_lr,
-            decay_steps=self.decay_step,
-            num_steps=self.stop_step,
-        )
-        with tf.Session(graph=g) as sess:
-            base_vals = [
-                sess.run(t_lr, feed_dict={global_step: step_id})
-                for step_id in range(self.stop_step)
-                if step_id % self.decay_step != 0
-            ]
-        my_vals = [
-            my_lr.value(step_id)
-            for step_id in range(self.stop_step)
-            if step_id % self.decay_step != 0
-        ]
-        self.assertTrue(np.allclose(base_vals, my_vals))
-        tf.reset_default_graph()
-
-    def decay_rate_pd(self):
-        my_lr = LearningRateExp(
-            start_lr=self.start_lr,
-            stop_lr=self.stop_lr,
-            decay_steps=self.decay_step,
-            num_steps=self.stop_step,
-        )
-
-        # Use the auto-adjusted decay_steps from my_lr for consistency
-        actual_decay_steps = my_lr.decay_steps
-        decay_rate = np.exp(
-            np.log(self.stop_lr / self.start_lr) / (self.stop_step / actual_decay_steps)
-        )
-        my_lr_decay = LearningRateExp(
-            start_lr=self.start_lr,
-            stop_lr=1e-10,
-            decay_steps=actual_decay_steps,
-            num_steps=self.stop_step,
-            decay_rate=decay_rate,
-        )
-        min_lr = 1e-5
-        my_lr_decay_trunc = LearningRateExp(
-            start_lr=self.start_lr,
-            stop_lr=min_lr,
-            decay_steps=actual_decay_steps,
-            num_steps=self.stop_step,
-            decay_rate=decay_rate,
-        )
-        my_vals = [
-            my_lr.value(step_id)
-            for step_id in range(self.stop_step)
-            if step_id % actual_decay_steps != 0
-        ]
-        my_vals_decay = [
-            my_lr_decay.value(step_id)
-            for step_id in range(self.stop_step)
-            if step_id % actual_decay_steps != 0
-        ]
-        my_vals_decay_trunc = [
-            my_lr_decay_trunc.value(step_id)
-            for step_id in range(self.stop_step)
-            if step_id % actual_decay_steps != 0
-        ]
-        self.assertTrue(np.allclose(my_vals_decay, my_vals))
-        self.assertTrue(
-            np.allclose(my_vals_decay_trunc, np.clip(my_vals, a_min=min_lr, a_max=None))
-        )
-
-
-class TestLearningRateWSD(unittest.TestCase):
-    def test_basic_curve(self):
-        start_lr = 1.0
-        stop_lr = 0.1
-        stop_steps = 10
-        lr = LearningRateWSD(
-            start_lr=start_lr,
-            stop_lr=stop_lr,
-            num_steps=stop_steps,
-            decay_phase_ratio=0.2,
-        )
-
-        expected_mid = 1.0 / (0.5 / stop_lr + 0.5 / start_lr)
-        self.assertTrue(np.allclose(lr.value(0), start_lr))
-        self.assertTrue(np.allclose(lr.value(8), start_lr))
-        self.assertTrue(np.allclose(lr.value(9), expected_mid))
-        self.assertTrue(np.allclose(lr.value(stop_steps), stop_lr))
-        self.assertTrue(np.allclose(lr.value(stop_steps + 5), stop_lr))
-
-    def test_cosine_curve(self):
-        start_lr = 1.0
-        stop_lr = 0.1
-        stop_steps = 10
-        lr = LearningRateWSD(
-            start_lr=start_lr,
-            stop_lr=stop_lr,
-            num_steps=stop_steps,
-            decay_phase_ratio=0.2,
-            decay_type="cosine",
-        )
-
-        expected_mid = stop_lr + (start_lr - stop_lr) * 0.5
-        self.assertTrue(np.allclose(lr.value(0), start_lr))
-        self.assertTrue(np.allclose(lr.value(8), start_lr))
-        self.assertTrue(np.allclose(lr.value(9), expected_mid))
-        self.assertTrue(np.allclose(lr.value(stop_steps), stop_lr))
-
-    def test_linear_curve(self):
-        start_lr = 1.0
-        stop_lr = 0.1
-        stop_steps = 10
-        lr = LearningRateWSD(
-            start_lr=start_lr,
-            stop_lr=stop_lr,
-            num_steps=stop_steps,
-            decay_phase_ratio=0.2,
-            decay_type="linear",
-        )
-
-        expected_mid = start_lr + (stop_lr - start_lr) * 0.5
-        self.assertTrue(np.allclose(lr.value(0), start_lr))
-        self.assertTrue(np.allclose(lr.value(8), start_lr))
-        self.assertTrue(np.allclose(lr.value(9), expected_mid))
-        self.assertTrue(np.allclose(lr.value(stop_steps), stop_lr))
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/source/tests/pd/test_sampler.py b/source/tests/pd/test_sampler.py
index e920660704..37a23fed06 100644
--- a/source/tests/pd/test_sampler.py
+++ b/source/tests/pd/test_sampler.py
@@ -2,9 +2,15 @@
 import json
 import os
 import unittest
+from collections.abc import (
+    Iterator,
+)
 from pathlib import (
     Path,
 )
+from threading import (
+    Event,
+)
 
 import numpy as np
 import paddle
@@ -17,6 +23,8 @@
     expand_sys_str,
 )
 from deepmd.pd.utils.dataloader import (
+    QUEUESIZE,
+    BufferedIterator,
     DpLoaderSet,
     get_weighted_sampler,
 )
@@ -28,6 +36,44 @@
 CUR_DIR = os.path.dirname(__file__)
 
 
+class TestBufferedIterator(unittest.TestCase):
+    def test_close_stops_production_with_a_full_queue(self) -> None:
+        full = Event()
+        extra_read = Event()
+
+        class Source:
+            def __len__(self) -> int:
+                return 10000
+
+            def __iter__(self) -> Iterator[int]:
+                for index in range(len(self)):
+                    if index == QUEUESIZE + 1:
+                        full.set()
+                    if index == QUEUESIZE + 2:
+                        extra_read.set()
+                    yield index
+
+        buffered = BufferedIterator(Source())
+        self.addCleanup(buffered.close)
+        self.assertEqual(next(buffered), 0)
+        self.assertTrue(full.wait(timeout=5))
+        self.assertTrue(buffered._queue.full())
+        producer = buffered._consumer
+        buffered.close()
+        self.assertFalse(producer.is_alive())
+        self.assertFalse(extra_read.is_set())
+        self.assertTrue(buffered._queue.empty())
+        with self.assertRaises(StopIteration):
+            next(buffered)
+        buffered.close()
+
+    def test_exhaustion_releases_the_producer(self) -> None:
+        buffered = BufferedIterator(range(3))
+        self.assertEqual(list(buffered), [0, 1, 2])
+        self.assertIsNone(buffered._consumer)
+        self.assertIsNone(buffered._iterable)
+
+
 class TestSampler(unittest.TestCase):
     def setUp(self):
         with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin:
diff --git a/source/tests/pd/test_update_sel.py b/source/tests/pd/test_update_sel.py
index 10342357c6..04f98d6ba1 100644
--- a/source/tests/pd/test_update_sel.py
+++ b/source/tests/pd/test_update_sel.py
@@ -144,33 +144,6 @@ def test_update_sel_atten_int(self, sel_mock) -> None:
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    @patch("deepmd.pd.utils.update_sel.UpdateSel.get_nbor_stat")
-    def test_update_sel_atten_list(self, sel_mock) -> None:
-        sel_mock.return_value = self.mock_min_nbor_dist, [25]
-
-        jdata = {
-            "model": {
-                "descriptor": {
-                    "type": "se_atten",
-                    "sel": 30,
-                    "rcut": 6,
-                }
-            },
-            "training": {"training_data": {}},
-        }
-        expected_out = {
-            "model": {
-                "descriptor": {
-                    "type": "se_atten",
-                    "sel": 30,
-                    "rcut": 6,
-                }
-            },
-            "training": {"training_data": {}},
-        }
-        jdata = update_sel(jdata)
-        self.assertEqual(jdata, expected_out)
-
     @patch("deepmd.pd.utils.update_sel.UpdateSel.get_nbor_stat")
     def test_update_sel_dpa2_auto(self, sel_mock) -> None:
         sel_mock.return_value = self.mock_min_nbor_dist, [25]
diff --git a/source/tests/pt/model/test_atomic_model_global_stat.py b/source/tests/pt/model/test_atomic_model_global_stat.py
index a4bca815b9..78c877bce9 100644
--- a/source/tests/pt/model/test_atomic_model_global_stat.py
+++ b/source/tests/pt/model/test_atomic_model_global_stat.py
@@ -7,6 +7,9 @@
 from typing import (
     NoReturn,
 )
+from unittest.mock import (
+    patch,
+)
 
 import h5py
 import numpy as np
@@ -386,6 +389,42 @@ def cvt_ret(x):
         out_bias, _ = md0._fetch_out_stat(["foo", "bar"])
         np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias)
         np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias)
+
+        # Standalone change-bias explicitly ignores the checkpoint preset.
+        adjusted_bias = foo_bias + 7.0
+        md0._store_out_stat({"foo": to_torch_tensor(adjusted_bias)})
+        forward = md0._get_forward_wrapper_func()
+        seen_biases = []
+
+        def check_forward(*args, **kwargs):
+            bias, _ = md0._fetch_out_stat(["foo"])
+            seen_biases.append(to_numpy_array(bias["foo"]).copy())
+            return forward(*args, **kwargs)
+
+        with patch.object(md0, "_get_forward_wrapper_func", return_value=check_forward):
+            BaseAtomicModel.change_out_bias(
+                md0,
+                self.merged_output_stat,
+                bias_adjust_mode="change-by-statistic",
+                preset_out_bias={},
+            )
+        self.assertTrue(seen_biases)
+        for seen_bias in seen_biases:
+            np.testing.assert_almost_equal(seen_bias, adjusted_bias)
+        out_bias, _ = md0._fetch_out_stat(["foo", "bar"])
+        np.testing.assert_almost_equal(
+            to_numpy_array(out_bias["foo"]), np.array([2.0, -5.0]).reshape(2, 1)
+        )
+        assert not np.allclose(to_numpy_array(out_bias["bar"]), bar_bias)
+
+        # The override does not remove the model's configured preset.
+        np.testing.assert_almost_equal(md0.preset_out_bias["foo"], foo_bias)
+        BaseAtomicModel.change_out_bias(
+            md0, self.merged_output_stat, bias_adjust_mode="change-by-statistic"
+        )
+        out_bias, _ = md0._fetch_out_stat(["foo", "bar"])
+        np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias)
+        np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias)
         # 4. a preset that leaves an observed type unassigned is rejected
         md1 = DPAtomicModel(
             ds,
@@ -659,6 +698,43 @@ def test_nonfinite_assigned_bias_from_checkpoint(self) -> None:
                     to_numpy_array(loaded.get_out_bias()).reshape(-1), bias, atol=1e-10
                 )
 
+    def test_empty_override_recovers_nonfinite_checkpoint_bias(self) -> None:
+        for value in (np.nan, np.inf, -np.inf):
+            with self.subTest(value=value):
+                loaded = EnergyModel.deserialize(self.model.serialize())
+                loaded.atomic_model.preset_out_bias = {
+                    "energy": [[-10.0], [-13.6], [-3.0]]
+                }
+                loaded.atomic_model.out_bias[0, :, 0] = to_torch_tensor(
+                    np.array([value, 7.0, value])
+                )
+                loaded = EnergyModel.deserialize(loaded.serialize())
+                forward = loaded.predict_atomic_outputs_for_stat
+                seen = []
+
+                def check_forward(*args, **kwargs):
+                    seen.append(
+                        to_numpy_array(loaded.get_out_bias()).reshape(-1).copy()
+                    )
+                    return forward(*args, **kwargs)
+
+                with patch.object(
+                    loaded, "predict_atomic_outputs_for_stat", check_forward
+                ):
+                    loaded.change_out_bias(self.sampled, preset_out_bias={})
+                self.assertTrue(seen)
+                for baseline in seen:
+                    np.testing.assert_allclose(baseline, [-10.0, 7.0, -3.0])
+                bias = to_numpy_array(loaded.get_out_bias()).reshape(-1)
+                self.assertTrue(np.isfinite(bias).all())
+                self.assertAlmostEqual(bias[2], -3.0)
+                prediction = loaded(self.sampled[0]["coord"], self.sampled[0]["atype"])
+                np.testing.assert_allclose(
+                    to_numpy_array(prediction["energy"]),
+                    to_numpy_array(self.sampled[0]["energy"]),
+                    atol=1e-5,
+                )
+
     def test_change_type_map_remaps_preset(self) -> None:
         self.model.change_out_bias(self.sampled, bias_adjust_mode="set-by-statistic")
         # H is dropped, O keeps its preset, B and C are unassigned
diff --git a/source/tests/pt/model/test_force_grad.py b/source/tests/pt/model/test_force_grad.py
deleted file mode 100644
index 599f0d1830..0000000000
--- a/source/tests/pt/model/test_force_grad.py
+++ /dev/null
@@ -1,108 +0,0 @@
-# SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
-import json
-import unittest
-from pathlib import (
-    Path,
-)
-
-import numpy as np
-import torch
-
-from deepmd.pt.model.model import (
-    get_model,
-)
-from deepmd.pt.utils import (
-    env,
-)
-from deepmd.utils.data import (
-    DeepmdData,
-)
-
-from ...seed import (
-    GLOBAL_SEED,
-)
-
-
-class CheckSymmetry(DeepmdData):
-    def __init__(
-        self,
-        sys_path: str,
-        type_map: list[str] | None = None,
-    ) -> None:
-        super().__init__(sys_path=sys_path, type_map=type_map)
-        self.add("energy", 1, atomic=False, must=False, high_prec=True)
-        self.add("force", 3, atomic=True, must=False, high_prec=False)
-        self.add("virial", 9, atomic=False, must=False, high_prec=False)
-
-    def get_disturb(self, index, atom_index, axis_index, delta):
-        for i in range(
-            len(self.dirs) + 1
-        ):  # note: if different sets can be merged, prefix sum is unused to calculate
-            if index < self.prefix_sum[i]:
-                break
-        frames = self._load_set(self.dirs[i - 1])
-        tmp = copy.deepcopy(frames["coord"].reshape(self.nframes, -1, 3))
-        tmp[:, atom_index, axis_index] += delta
-        frames["coord"] = tmp
-        frame = self._get_subdata(frames, index - self.prefix_sum[i - 1])
-        frame = self.reformat_data_torch(frame)
-        return frame
-
-
-def get_data(batch):
-    inputs = {}
-    for key in ["coord", "atype", "box"]:
-        inputs[key] = batch[key].unsqueeze(0).to(env.DEVICE)
-    return inputs
-
-
-class TestForceGrad(unittest.TestCase):
-    def setUp(self) -> None:
-        with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin:
-            self.config = json.load(fin)
-        data_file = [str(Path(__file__).parent / "water/data/data_0")]
-        self.config["training"]["training_data"]["systems"] = data_file
-        self.config["training"]["validation_data"]["systems"] = data_file
-        self.system_index = 0
-        self.batch_index = 0
-        self.get_dataset(self.system_index, self.batch_index)
-        self.get_model()
-
-    def get_model(self) -> None:
-        self.model = get_model(self.config["model"]).to(env.DEVICE)
-
-    def get_dataset(self, system_index=0, batch_index=0) -> None:
-        systems = self.config["training"]["training_data"]["systems"]
-        rcut = self.config["model"]["descriptor"]["rcut"]
-        sel = self.config["model"]["descriptor"]["sel"]
-        sec = torch.cumsum(torch.tensor(sel), dim=0)
-        type_map = self.config["model"]["type_map"]
-        self.dpdatasystem = CheckSymmetry(
-            sys_path=systems[system_index], type_map=type_map
-        )
-        self.origin_batch = self.dpdatasystem.get_item_torch(batch_index)
-
-    @unittest.skip("it can be replaced by autodiff")
-    def test_force_grad(self, threshold=1e-2, delta0=1e-6, seed=20) -> None:
-        rng = np.random.default_rng(GLOBAL_SEED)
-        result0 = self.model(**get_data(self.origin_batch))
-        np.random.default_rng(seed)
-        errors = np.zeros((self.dpdatasystem.natoms, 3))
-        for atom_index in range(self.dpdatasystem.natoms):
-            for axis_index in range(3):
-                delta = rng.random() * delta0
-                disturb_batch = self.dpdatasystem.get_disturb(
-                    self.batch_index, atom_index, axis_index, delta
-                )
-                disturb_result = self.model(**get_data(disturb_batch))
-                disturb_force = -(disturb_result["energy"] - result0["energy"]) / delta
-                disturb_error = (
-                    result0["force"][0, atom_index, axis_index] - disturb_force
-                )
-                errors[atom_index, axis_index] = disturb_error.detach().cpu().numpy()
-        self.assertTrue(np.abs(errors).max() < threshold, msg=str(np.abs(errors).max()))
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/source/tests/pt/model/test_get_model_bridging.py b/source/tests/pt/model/test_get_model_bridging.py
index 614952641f..0089e69972 100644
--- a/source/tests/pt/model/test_get_model_bridging.py
+++ b/source/tests/pt/model/test_get_model_bridging.py
@@ -394,24 +394,14 @@ def test_canonical_conflicting_top_level_option_raises() -> None:
         get_model(cfg)
 
 
-@pytest.mark.parametrize(
-    "scheme",
-    [
-        "native",  # spin as an equivariant descriptor feature
-        "deepspin",  # classical virtual-atom representation
-    ],
-)
-def test_canonical_spin_rejects_mismatched_pair_exclusions(scheme: str) -> None:
-    """Both spin routes must fail fast on a pair-exclusion mismatch like
-    the no-spin route, not silently overwrite the descriptor's exclusions.
-    """
+def test_canonical_spin_rejects_mismatched_pair_exclusions() -> None:
+    """Spin models reject conflicting pair exclusions."""
     cfg = _canonical_config()
     cfg["pair_exclude_types"] = [[0, 0]]
     cfg["models"][0]["descriptor"]["exclude_types"] = [[0, 1]]
     cfg["spin"] = {
-        "scheme": scheme,
+        "scheme": "native",
         "use_spin": [True, False],
-        "virtual_scale": 0.3,
     }
     with pytest.raises(ValueError, match="must match"):
         get_model(cfg)
@@ -432,3 +422,15 @@ def test_normalized_canonical_top_level_option_survives_child_defaults() -> None
     assert cfg["models"][0]["data_stat_protect"] == 0.01  # injected default
     model = get_model(cfg)
     assert model.atomic_model.data_stat_protect == 0.123
+
+
+@pytest.mark.parametrize("canonical", [False, True])
+@pytest.mark.parametrize("scheme", [None, "invalid"])
+def test_spin_requires_native_scheme(canonical: bool, scheme: str | None) -> None:
+    """Both public model spellings require an explicit native spin scheme."""
+    cfg = _canonical_config() if canonical else _sugar_config()
+    cfg["spin"] = {"use_spin": [True, False]}
+    if scheme is not None:
+        cfg["spin"]["scheme"] = scheme
+    with pytest.raises(ValueError, match=r"require model\.spin\.scheme='native'"):
+        get_model(cfg)
diff --git a/source/tests/pt/model/test_nlist_backend.py b/source/tests/pt/model/test_nlist_backend.py
index d38262dacd..7b757efd65 100644
--- a/source/tests/pt/model/test_nlist_backend.py
+++ b/source/tests/pt/model/test_nlist_backend.py
@@ -5,11 +5,16 @@
 ``forward_common_lower`` when an O(N) neighbor-list strategy is selected.  Each
 strategy (``vesin``, ``nv``) must give results identical to the native dense
 builder, and the ``nlist_backend`` choice must dispatch / validate correctly.
-Strategy tests are skipped when the backend (or, for ``nv``, a CUDA device) is
-unavailable.
+Strategy tests are skipped when the backend is unavailable.
 """
 
 import copy
+from pathlib import (
+    Path,
+)
+from unittest import (
+    mock,
+)
 
 import numpy as np
 import pytest
@@ -24,6 +29,9 @@
 from deepmd.pt.train.wrapper import (
     ModelWrapper,
 )
+from deepmd.pt.utils.env import (
+    DEVICE,
+)
 from deepmd.pt.utils.nv_nlist import (
     NvNeighborList,
     is_nv_available,
@@ -33,14 +41,17 @@
     is_vesin_torch_available,
 )
 
+from ...common.sezm_config import (
+    _tiny_sezm_model_params,
+)
+
 # Each O(N) strategy: (backend name, builder class, availability skip mark).
 _BACKEND_MARKS = {
     "vesin": pytest.mark.skipif(
         not is_vesin_torch_available(), reason="vesin.torch is not installed"
     ),
     "nv": pytest.mark.skipif(
-        not (is_nv_available() and torch.cuda.is_available()),
-        reason="nvalchemiops CUDA neighbor list unavailable",
+        not is_nv_available(), reason="nvalchemiops neighbor list unavailable"
     ),
 }
 _BUILDER_CLS = {"vesin": VesinNeighborList, "nv": NvNeighborList}
@@ -143,6 +154,31 @@ def test_invalid_backend_raises(pt_files) -> None:
         DeepPot(pt_files["se_e2_a"], nlist_backend="bogus")
 
 
+@pytest.mark.parametrize(
+    ("method", "backend", "message"),
+    [
+        ("dense", "auto", "PT SeZM/DPA4 supports neighbor_graph_method"),
+        ("vesin", "native", "Select only one"),
+    ],
+)
+def test_self_built_neighbor_options_rejected(
+    tmp_path: Path, method: str, backend: str, message: str
+) -> None:
+    path = str(tmp_path / "sezm.pt")
+    _save_pt(_tiny_sezm_model_params(), path)
+    with pytest.raises(ValueError, match=message):
+        DeepPot(path, neighbor_graph_method=method, nlist_backend=backend)
+
+
+def test_neighbor_graph_method_rejected_for_dense_model(
+    pt_files: dict[str, str],
+) -> None:
+    with pytest.raises(
+        ValueError, match="neighbor_graph_method applies to PT SeZM/DPA4"
+    ):
+        DeepPot(pt_files["se_e2_a"], neighbor_graph_method="vesin")
+
+
 def test_native_uses_no_strategy(pt_files) -> None:
     dp = DeepPot(pt_files["se_e2_a"], nlist_backend="native")
     assert dp.deep_eval._nlist_builder is None
@@ -155,10 +191,57 @@ def test_explicit_backend_selects_builder(pt_files, backend: str) -> None:
 
 
 @_BACKEND_MARKS["vesin"]
-def test_auto_prefers_vesin(pt_files) -> None:
-    # auto picks the first available O(N) builder; vesin is preferred.
-    builder = DeepPot(pt_files["se_e2_a"]).deep_eval._nlist_builder
-    assert isinstance(builder, VesinNeighborList)
+@pytest.mark.parametrize("nframes", [1, 3])
+def test_auto_selects_backend_from_frame_count(
+    pt_files: dict[str, str], nframes: int
+) -> None:
+    coords, atype, boxes = _multiframe_system(nframes)
+    dp = DeepPot(pt_files["se_e2_a"], auto_batch_size=False)
+    with (
+        mock.patch.object(
+            VesinNeighborList,
+            "build",
+            autospec=True,
+            side_effect=VesinNeighborList.build,
+        ) as vesin_build,
+        mock.patch.object(
+            NvNeighborList, "build", autospec=True, side_effect=NvNeighborList.build
+        ) as nv_build,
+    ):
+        dp.eval(coords, boxes, atype)
+
+    if nframes > 1 and is_nv_available():
+        nv_build.assert_called_once()
+        vesin_build.assert_not_called()
+    else:
+        vesin_build.assert_called_once()
+        nv_build.assert_not_called()
+
+
+@pytest.mark.parametrize("backend", STRATEGIES)
+@pytest.mark.parametrize("option", ["nlist_backend", "neighbor_graph_method"])
+def test_self_built_model_uses_inference_strategy(
+    tmp_path: Path, backend: str, option: str
+) -> None:
+    path = str(tmp_path / "sezm.pt")
+    _save_pt(_tiny_sezm_model_params(), path)
+    dp = DeepPot(path, auto_batch_size=False, **{option: backend})
+    model = dp.deep_eval.dp.model["Default"]
+    builder = model._inference_neighbor_list
+    assert isinstance(builder, _BUILDER_CLS[backend])
+    coords = np.array([[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [0.0, 1.1, 0.0]]])
+    atype = np.array([0, 1, 0])
+    with mock.patch.object(builder, "build", wraps=builder.build) as build:
+        dp.eval(coords, None, atype)
+        build.assert_called_once()
+        assert model._get_neighbor_builder() is builder
+        model.train()
+        model.build_neighbor_list(
+            torch.as_tensor(coords, dtype=torch.float64, device=DEVICE),
+            torch.as_tensor(atype.reshape(1, -1), device=DEVICE),
+            None,
+        )
+        build.assert_called_once()
 
 
 def test_self_built_model_forces_native(pt_files, monkeypatch) -> None:
diff --git a/source/tests/pt/model/test_sezm_export.py b/source/tests/pt/model/test_sezm_export.py
index 04a920d83f..037bc2d4ba 100644
--- a/source/tests/pt/model/test_sezm_export.py
+++ b/source/tests/pt/model/test_sezm_export.py
@@ -46,6 +46,11 @@
     ensure_comm_registered,
 )
 
+from ...common.sezm_config import (
+    _tiny_sezm_model_params,
+    _tiny_sezm_native_spin_model_params,
+)
+
 _COMM_KEYS = (
     "send_list",
     "send_proc",
@@ -133,76 +138,6 @@ def _eager_parallel_forward(
 )
 
 
-def _tiny_sezm_model_params() -> dict:
-    """Minimal fp64 SeZM config for self-contained export tests.
-
-    ``precision="float64"`` is what unlocks the ``rtol=1e-10, atol=1e-10``
-    parity pt_expt enforces; fp32 accumulation alone drifts in the 1e-6
-    range.  All other knobs are tuned to keep ``make_fx`` tracing time
-    in the low-single-digit seconds.
-    """
-    return {
-        "type": "SeZM",
-        "type_map": ["A", "B"],
-        "descriptor": {
-            "type": "SeZM",
-            "sel": [2, 2],
-            "rcut": 3.0,
-            "channels": 4,
-            "n_focus": 1,
-            "n_radial": 3,
-            "radial_mlp": [6],
-            "use_env_seed": True,
-            "l_schedule": [1, 0],
-            "mmax": 1,
-            "so2_norm": False,
-            "so2_layers": 1,
-            "n_atten_head": 1,
-            "sandwich_norm": [True, False, True, False],
-            "ffn_neurons": 8,
-            "ffn_blocks": 1,
-            "s2_activation": [False, True],
-            "mlp_bias": False,
-            "layer_scale": False,
-            "use_amp": False,
-            "activation_function": "silu",
-            "glu_activation": True,
-            "precision": "float64",
-            "seed": 7,
-        },
-        "fitting_net": {
-            "neuron": [8],
-            "activation_function": "silu",
-            "precision": "float64",
-            "seed": 7,
-        },
-        "use_compile": False,
-    }
-
-
-def _tiny_sezm_spin_model_params() -> dict:
-    """Minimal fp64 SeZM spin config for freeze routing tests."""
-    params = copy.deepcopy(_tiny_sezm_model_params())
-    params["type_map"] = ["O", "H"]
-    params["spin"] = {
-        "use_spin": [True, False],
-        "virtual_scale": 0.2,
-        "scheme": "deepspin",
-    }
-    return params
-
-
-def _tiny_sezm_native_spin_model_params() -> dict:
-    """Minimal fp64 native-spin SeZM config for freeze routing tests."""
-    params = copy.deepcopy(_tiny_sezm_model_params())
-    params["type_map"] = ["O", "H"]
-    params["spin"] = {
-        "use_spin": [True, False],
-        "scheme": "native",
-    }
-    return params
-
-
 def _build_tiny_sezm_model() -> torch.nn.Module:
     """Fresh tiny SeZM model on CPU, in eval mode."""
     model = get_model(_tiny_sezm_model_params())
@@ -1084,47 +1019,6 @@ def fake_compile(_exported: torch.export.ExportedProgram, package_path: str):
 
         self.assertEqual(model_def["type"], "dpa4")
 
-    @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
-    def test_freeze_accepts_spin_checkpoint_metadata(self) -> None:
-        """SeZM spin checkpoints should export a spin-compatible pt2 contract."""
-
-        def fake_compile(_exported: torch.export.ExportedProgram, package_path: str):
-            with zipfile.ZipFile(package_path, "w") as zf:
-                zf.writestr("model/data.pkl", b"")
-
-        with tempfile.TemporaryDirectory() as tmp:
-            tmp_path = Path(tmp)
-            params = _tiny_sezm_spin_model_params()
-            ckpt_path = _write_tiny_sezm_checkpoint(tmp_path, params)
-            out = tmp_path / "spin.pt2"
-
-            with mock.patch(
-                "torch._inductor.aoti_compile_and_package",
-                side_effect=fake_compile,
-            ):
-                freeze_sezm_to_pt2(str(ckpt_path), str(out), device=_CPU)
-
-            with zipfile.ZipFile(str(out), "r") as zf:
-                names = zf.namelist()
-                metadata = json.loads(
-                    zf.read("model/extra/metadata.json").decode("utf-8")
-                )
-
-        self.assertTrue(metadata["is_spin"])
-        self.assertEqual(metadata["lower_input_kind"], "nlist")
-        self.assertEqual(metadata["type_map"], params["type_map"])
-        self.assertEqual(metadata["ntypes"], len(params["type_map"]))
-        self.assertEqual(metadata["dim_chg_spin"], 0)
-        self.assertIsNone(metadata["default_chg_spin"])
-        self.assertEqual(metadata["use_spin"], params["spin"]["use_spin"])
-        self.assertEqual(metadata["ntypes_spin"], 1)
-        self.assertIn("energy_derv_r_mag", metadata["output_keys"])
-        self.assertIn("energy_derv_c_redu", metadata["output_keys"])
-        # Spin uses the nlist lower interface; the edge-based with-comm artifact
-        # does not apply, so multi-rank inference fails fast in C++.
-        self.assertFalse(metadata["has_comm_artifact"])
-        self.assertNotIn("model/extra/forward_lower_with_comm.pt2", names)
-
     @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
     def test_freeze_accepts_native_spin_checkpoint_metadata(self) -> None:
         """Native-spin checkpoints export the energy edge contract plus spin.
@@ -1132,8 +1026,7 @@ def test_freeze_accepts_native_spin_checkpoint_metadata(self) -> None:
         The native scheme reuses the ``edge_vec`` lower ABI; the per-local-atom
         spin is the only extra input, so the C++ backend builds the edge schema
         exactly as for a non-spin model. The magnetic force and spin mask are
-        still emitted, and the type map / ntypes stay at the real-system sizes
-        (no virtual atoms).
+        still emitted, and the type map / ntypes stay at the real-system sizes.
         """
 
         def fake_compile(_exported: torch.export.ExportedProgram, package_path: str):
@@ -1159,12 +1052,13 @@ def fake_compile(_exported: torch.export.ExportedProgram, package_path: str):
                 )
 
         self.assertTrue(metadata["is_spin"])
-        # Native spin shares the energy edge ABI; only the deepspin scheme keeps
-        # the nlist contract.
+        # Native spin shares the energy edge ABI.
         self.assertEqual(metadata["lower_input_kind"], "edge_vec")
-        # Native spin keeps the real-system type map and count (no virtual atoms).
+        # Native spin keeps the real-system type map and count.
         self.assertEqual(metadata["type_map"], params["type_map"])
         self.assertEqual(metadata["ntypes"], len(params["type_map"]))
+        self.assertEqual(metadata["sel"], params["descriptor"]["sel"])
+        self.assertEqual(metadata["nnei"], sum(params["descriptor"]["sel"]))
         self.assertEqual(metadata["use_spin"], params["spin"]["use_spin"])
         self.assertEqual(metadata["ntypes_spin"], 1)
         # The magnetic force and spin mask are still exported.
diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py
index 18ca49b8bf..7bf68d9e00 100644
--- a/source/tests/pt/model/test_sezm_model.py
+++ b/source/tests/pt/model/test_sezm_model.py
@@ -476,7 +476,7 @@ def test_eval_compile_first_frame_nloc_matches_trace_edge_count(self) -> None:
         # The compiled eval path must trace, lower and run without tripping
         # ``assert_size_stride`` on the edge tensors.
         out_cmp = model_cmp(coord, atype, box=box)
-        self.assertIn((False, False), model_cmp.compiled_core_compute_cache)
+        self.assertIn(False, model_cmp.compiled_core_compute_cache)
         _assert_close_with_strict_warning(
             out_dyn["energy"],
             out_cmp["energy"],
@@ -505,9 +505,9 @@ def test_compile_cache_slots_and_eval_shape_change(self) -> None:
             model_cmp = get_sezm_model(self._build_model_params(use_compile=True))
         model_cmp.load_state_dict(model_dyn.state_dict())
 
-        # Compile cache key is (training, has_coord_corr).
-        train_key = (True, False)
-        eval_key = (False, False)
+        # Compile cache key is the training mode.
+        train_key = True
+        eval_key = False
 
         # === Step 2. Train-mode forward fills the training slot. ===
         model_cmp.train()
@@ -1081,13 +1081,13 @@ def _build_wrapper(use_compile: bool) -> ModelWrapper:
         # the module-level ``_SEZM_COMPILE_CACHE``.  This avoids the
         # N x compile-cache OOM / N DDP graph boundary cost on multitask
         # runs.  Step 2 ran every branch in training mode with the default
-        # ``do_atomic_virial=False`` and no coordinate correction, so each
+        # ``do_atomic_virial=False``, so each
         # per-branch cache should hold exactly that one slot, and the
         # compiled callable at that slot must be the *same* object. ===
         cache1 = wrapper_cmp.model["water_1"].compiled_core_compute_cache
         cache2 = wrapper_cmp.model["water_2"].compiled_core_compute_cache
         self.assertIsNot(cache1, cache2)
-        train_key = (True, False)
+        train_key = True
         self.assertIn(train_key, cache1)
         self.assertIn(train_key, cache2)
         c1 = cache1[train_key]
@@ -1304,7 +1304,7 @@ def test_compile_matches_eager_and_backpropagates(self) -> None:
                 rtol=forward_tol,
                 msg=f"compiled property mismatch at {key}",
             )
-        self.assertIn((True, False), compiled.compiled_core_compute_cache)
+        self.assertIn(True, compiled.compiled_core_compute_cache)
 
         # === Step 3. Compare parameter gradients ===
         probe = torch.tensor(
@@ -1696,7 +1696,7 @@ def test_atom_virial_sums_to_global_virial(self) -> None:
 
 
 class TestSeZMNativeSpinModel(unittest.TestCase):
-    """Validate the native (virtual-atom-free) spin SeZM model.
+    """Validate the native spin SeZM model.
 
     The spin vector enters the descriptor as an equivariant feature, so the
     magnetic force is the negative spin gradient of the energy. float64
@@ -2112,7 +2112,7 @@ def test_compile_matches_eager(self) -> None:
 
         out_e = model_eager(coord, atype, spin, box=box)
         out_c = model_cmp(coord, atype, spin, box=box)
-        self.assertIn((True, False), model_cmp.compiled_core_compute_cache)
+        self.assertIn(True, model_cmp.compiled_core_compute_cache)
         for key in ["energy", "force", "force_mag"]:
             _assert_close_with_strict_warning(
                 out_c[key],
diff --git a/source/tests/pt/model/test_sezm_spin_contract.py b/source/tests/pt/model/test_sezm_spin_contract.py
new file mode 100644
index 0000000000..62ca575892
--- /dev/null
+++ b/source/tests/pt/model/test_sezm_spin_contract.py
@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""DPA4 spin construction contracts for the PyTorch backend."""
+
+import copy
+
+import pytest
+
+from deepmd.pt.model.model import (
+    get_model,
+    get_standard_model,
+)
+from deepmd.pt.model.model.dp_linear_model import (
+    LinearEnergyModel,
+)
+from deepmd.pt.model.model.sezm_native_spin_model import (
+    SeZMNativeSpinModel,
+)
+from deepmd.pt.model.model.spin_model import (
+    SpinEnergyModel,
+)
+from deepmd.utils.spin import (
+    Spin,
+)
+
+
+def _config() -> dict:
+    return {
+        "type_map": ["Ni", "O"],
+        "descriptor": {
+            "type": "dpa4",
+            "rcut": 4.0,
+            "sel": 8,
+            "channels": 4,
+            "n_radial": 3,
+            "lmax": 1,
+            "mmax": 1,
+            "n_blocks": 1,
+            "precision": "float64",
+            "seed": 7,
+            "random_gamma": False,
+        },
+        "fitting_net": {"type": "ener", "neuron": [4]},
+    }
+
+
+@pytest.mark.parametrize("hybrid", [False, True])
+@pytest.mark.parametrize("descriptor_type", ["dpa4", "DPA4", "sezm", "SeZM"])
+@pytest.mark.parametrize("model_type", [None, "standard"])
+@pytest.mark.parametrize("scheme", [None, "invalid"])
+def test_standard_spin_requires_compatible_backbone(
+    descriptor_type: str, model_type: str | None, scheme: str | None, hybrid: bool
+) -> None:
+    config = _config()
+    config["descriptor"]["type"] = descriptor_type
+    if hybrid:
+        config["descriptor"] = {"type": "hybrid", "list": [config["descriptor"]]}
+    if model_type is not None:
+        config["type"] = model_type
+    config["spin"] = {"use_spin": [True, False], "virtual_scale": 0.3}
+    if scheme is not None:
+        config["spin"]["scheme"] = scheme
+    with pytest.raises(ValueError, match=r"require.*native"):
+        get_model(config)
+
+
+def test_native_spin_construction() -> None:
+    config = _config()
+    config["spin"] = {"use_spin": [True, False], "scheme": "native"}
+    config["type"] = "dpa4"
+    config["fitting_net"]["type"] = "dpa4_ener"
+    original = copy.deepcopy(config)
+    model = get_model(config)
+    assert isinstance(model, SeZMNativeSpinModel)
+    assert model.get_type_map() == ["Ni", "O"]
+    assert model.get_sel() == [config["descriptor"]["sel"]]
+    assert model.get_nnei() == config["descriptor"]["sel"]
+    assert config == original
+
+
+@pytest.mark.parametrize(
+    "serialized, depth", [(False, 0), (False, 1), (False, 2), (True, 0)]
+)
+def test_spin_wrapper_requires_compatible_backbone(
+    serialized: bool, depth: int
+) -> None:
+    backbone = get_standard_model(_config())
+    for _ in range(depth):
+        backbone = LinearEnergyModel(
+            models=[backbone.atomic_model], type_map=backbone.get_type_map()
+        )
+    spin = Spin(use_spin=[True, False], virtual_scale=0.3)
+    with pytest.raises(ValueError, match="require a native spin model"):
+        if serialized:
+            SpinEnergyModel.deserialize(
+                {
+                    "type": "spin_ener",
+                    "backbone_model": backbone.serialize(),
+                    "spin": spin.serialize(),
+                }
+            )
+        else:
+            SpinEnergyModel(backbone_model=backbone, spin=spin)
+
+
+@pytest.mark.parametrize("serialized", [False, True])
+def test_spin_wrapper_checks_hybrid_descriptors(serialized: bool) -> None:
+    config = _config()
+    config["descriptor"] = {"type": "hybrid", "list": [config["descriptor"]]}
+    backbone = get_standard_model(config)
+    spin = Spin(use_spin=[True, False], virtual_scale=0.3)
+    with pytest.raises(ValueError, match="require a native spin model"):
+        if serialized:
+            SpinEnergyModel.deserialize(
+                {
+                    "type": "spin_ener",
+                    "backbone_model": backbone.serialize(),
+                    "spin": spin.serialize(),
+                }
+            )
+        else:
+            SpinEnergyModel(backbone_model=backbone, spin=spin)
diff --git a/source/tests/pt/model/test_sezm_spin_model.py b/source/tests/pt/model/test_sezm_spin_model.py
deleted file mode 100644
index ffb86609fb..0000000000
--- a/source/tests/pt/model/test_sezm_spin_model.py
+++ /dev/null
@@ -1,419 +0,0 @@
-# SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
-import json
-import os
-import tempfile
-import unittest
-import warnings
-from unittest import (
-    mock,
-)
-
-import torch
-from packaging.version import parse as parse_version
-
-from deepmd.pt.loss import (
-    EnergySpinLoss,
-)
-from deepmd.pt.model.model import (
-    get_model,
-)
-from deepmd.pt.model.model.sezm_spin_model import (
-    SeZMSpinModel,
-)
-from deepmd.pt.train.training import (
-    prepare_model_for_loss,
-)
-from deepmd.pt.utils import (
-    env,
-)
-from deepmd.pt.utils.compile_compat import (
-    SUPPORTED_COMPILE_TORCH,
-)
-from deepmd.pt.utils.nlist import (
-    extend_input_and_build_neighbor_list,
-)
-from deepmd.pt.utils.serialization import (
-    deserialize_to_file,
-)
-
-warnings.filterwarnings(
-    # Keep the compile-test warning summary focused on strict-tolerance drift.
-    # PyTorch's AOTAutograd cache emits an internal Python 3.14 deprecation
-    # warning that is unrelated to SeZM numerical correctness.
-    "ignore",
-    category=DeprecationWarning,
-    module=r"torch\._functorch\._aot_autograd\.autograd_cache",
-)
-
-# Keep compile-parity test gating aligned with the runtime allowlist in
-# ``deepmd.pt.utils.compile_compat``. Membership in the allowlist does not imply
-# that every release is installed in CI.
-_TORCH_VERSION = parse_version(torch.__version__)
-_SKIP_OFF_COMPILE_TORCH = (
-    _TORCH_VERSION.major,
-    _TORCH_VERSION.minor,
-) not in SUPPORTED_COMPILE_TORCH
-_SKIP_OFF_COMPILE_TORCH_REASON = (
-    "SeZM's torch.compile path is only supported on torch "
-    + ", ".join(f"{major}.{minor}.x" for major, minor in SUPPORTED_COMPILE_TORCH)
-    + f"; current torch is {torch.__version__}."
-)
-
-
-def _assert_close_with_strict_warning(
-    actual: torch.Tensor,
-    expected: torch.Tensor,
-    *,
-    strict_atol: float = 1.0e-6,
-    strict_rtol: float = 1.0e-6,
-    atol: float,
-    rtol: float,
-    msg: str,
-) -> None:
-    """Warn on strict compile drift, fail only outside relaxed tolerance."""
-    try:
-        torch.testing.assert_close(
-            actual,
-            expected,
-            atol=strict_atol,
-            rtol=strict_rtol,
-            msg=msg,
-        )
-    except AssertionError as err:
-        warnings.warn(
-            f"{msg} exceeds strict tolerance "
-            f"(atol={strict_atol:g}, rtol={strict_rtol:g}) but is checked "
-            f"against relaxed tolerance (atol={atol:g}, rtol={rtol:g}): {err}",
-            RuntimeWarning,
-            stacklevel=2,
-        )
-    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol, msg=msg)
-
-
-def reduce_tensor(
-    extended_tensor: torch.Tensor,
-    mapping: torch.Tensor,
-    nloc: int,
-) -> torch.Tensor:
-    """Reduce an extended tensor back to local atoms."""
-    nframes = extended_tensor.shape[0]
-    ext_dims = extended_tensor.shape[2:]
-    reduced_tensor = torch.zeros(
-        [nframes, nloc, *ext_dims],
-        dtype=extended_tensor.dtype,
-        device=extended_tensor.device,
-    )
-    mldims = list(mapping.shape)
-    mapping = mapping.view(mldims + [1] * len(ext_dims)).expand(
-        [-1] * len(mldims) + list(ext_dims)
-    )
-    return torch.scatter_reduce(
-        reduced_tensor,
-        1,
-        index=mapping,
-        src=extended_tensor,
-        reduce="sum",
-    )
-
-
-class TestSeZMSpinModel(unittest.TestCase):
-    """Test spin support for the SeZM PyTorch model."""
-
-    def setUp(self) -> None:
-        self.device = env.DEVICE
-        torch.manual_seed(2024)
-        self.coord = torch.tensor(
-            [
-                [
-                    [0.0, 0.0, 0.0],
-                    [1.0, 0.0, 0.0],
-                    [0.0, 1.0, 0.0],
-                ]
-            ],
-            dtype=torch.float64,
-            device=self.device,
-        )
-        self.atype = torch.tensor([[0, 1, 0]], dtype=torch.long, device=self.device)
-        self.spin = torch.tensor(
-            [
-                [
-                    [0.20, 0.10, 0.00],
-                    [0.30, 0.00, 0.10],
-                    [0.10, 0.20, 0.10],
-                ]
-            ],
-            dtype=torch.float64,
-            device=self.device,
-        )
-        self.box = torch.eye(3, dtype=torch.float64, device=self.device).reshape(1, 9)
-        self.box = self.box * 6.0
-
-    def _build_model_params(
-        self,
-        *,
-        use_compile: bool = False,
-        bridging_method: str = "none",
-    ) -> dict:
-        """Build a minimal deterministic SeZM spin model config."""
-        return {
-            "type": "SeZM",
-            "type_map": ["O", "H"],
-            "spin": {
-                "use_spin": [True, False],
-                "virtual_scale": 0.2,
-                "scheme": "deepspin",
-            },
-            "descriptor": {
-                "type": "SeZM",
-                "sel": [2, 2],
-                "rcut": 3.0,
-                "channels": 4,
-                "n_focus": 1,
-                "n_radial": 3,
-                "radial_mlp": [6],
-                "use_env_seed": True,
-                "random_gamma": False,
-                "l_schedule": [1, 0],
-                "mmax": 1,
-                "so2_norm": False,
-                "so2_layers": 1,
-                "n_atten_head": 1,
-                "sandwich_norm": [True, False, True, False],
-                "ffn_neurons": 8,
-                "ffn_blocks": 1,
-                "s2_activation": [False, True],
-                "mlp_bias": False,
-                "layer_scale": False,
-                "use_amp": False,
-                "activation_function": "silu",
-                "glu_activation": True,
-                "precision": "float32",
-                "seed": 7,
-            },
-            "fitting_net": {
-                "neuron": [8],
-                "activation_function": "silu",
-                "precision": "float32",
-                "seed": 7,
-            },
-            "bridging_method": bridging_method,
-            "bridging_r_inner": 0.8,
-            "bridging_r_outer": 1.2,
-            "use_compile": use_compile,
-        }
-
-    def test_adam_route_patterns_match_parameters(self) -> None:
-        """The native-spin SeZM model routes the same tensors as the plain model."""
-        model = get_model(self._build_model_params()).to(self.device)
-        patterns = sorted(model.adam_route_patterns())
-        self.assertEqual(
-            patterns,
-            [
-                "atomic_model.descriptor.env_seed_embedding.rbf_proj_layer1.",
-                "atomic_model.descriptor.radial_embedding.net.0.",
-            ],
-        )
-        names = [name for name, _ in model.named_parameters()]
-        for pattern in patterns:
-            self.assertTrue(any(pattern in name for name in names), pattern)
-
-    def test_factory_shapes_and_masks(self) -> None:
-        """Factory should build SeZMSpinModel with public real-type metadata."""
-        model = get_model(self._build_model_params()).to(self.device)
-
-        self.assertIsInstance(model, SeZMSpinModel)
-        self.assertTrue(model.has_spin())
-        self.assertEqual(model.get_type_map(), ["O", "H"])
-        self.assertEqual(model.get_sel(), [2, 2])
-
-        out = model(self.coord, self.atype, spin=self.spin, box=self.box)
-
-        self.assertEqual(out["energy"].shape, (1, 1))
-        self.assertEqual(out["atom_energy"].shape, (1, 3, 1))
-        self.assertEqual(out["force"].shape, (1, 3, 3))
-        self.assertEqual(out["force_mag"].shape, (1, 3, 3))
-        torch.testing.assert_close(
-            out["mask_mag"],
-            torch.tensor(
-                [[[True], [False], [True]]],
-                dtype=torch.bool,
-                device=self.device,
-            ),
-        )
-
-    def test_forward_lower_matches_forward(self) -> None:
-        """Lower spin interface should match the standard spin forward path."""
-        model = get_model(self._build_model_params()).to(self.device)
-        model.eval()
-        out = model(self.coord, self.atype, spin=self.spin, box=self.box)
-        extended_coord, extended_atype, mapping, nlist = (
-            extend_input_and_build_neighbor_list(
-                self.coord,
-                self.atype,
-                model.get_rcut(),
-                model.get_sel(),
-                mixed_types=model.mixed_types(),
-                box=self.box,
-            )
-        )
-        extended_spin = torch.gather(
-            self.spin,
-            1,
-            mapping.unsqueeze(-1).expand(-1, -1, 3),
-        )
-
-        out_lower = model.forward_lower(
-            extended_coord,
-            extended_atype,
-            extended_spin,
-            nlist,
-            mapping=mapping,
-        )
-
-        torch.testing.assert_close(out_lower["energy"], out["energy"])
-        torch.testing.assert_close(out_lower["atom_energy"], out["atom_energy"])
-        reduced_force = reduce_tensor(out_lower["extended_force"], mapping, nloc=3)
-        reduced_force_mag = reduce_tensor(
-            out_lower["extended_force_mag"], mapping, nloc=3
-        )
-        torch.testing.assert_close(reduced_force, out["force"])
-        torch.testing.assert_close(reduced_force_mag, out["force_mag"])
-
-    def test_serialize_deserialize_consistency(self) -> None:
-        """Serialized SeZMSpinModel should restore the same predictions."""
-        model = get_model(self._build_model_params()).to(self.device)
-        restored = SeZMSpinModel.deserialize(model.serialize()).to(self.device)
-
-        out = model(self.coord, self.atype, spin=self.spin, box=self.box)
-        restored_out = restored(self.coord, self.atype, spin=self.spin, box=self.box)
-
-        self.assertEqual(restored.get_type_map(), ["O", "H"])
-        self.assertEqual(restored.get_sel(), [2, 2])
-        for key, value in out.items():
-            torch.testing.assert_close(restored_out[key], value)
-
-    def test_deserialize_to_file_uses_spin_model(self) -> None:
-        """File deserialization should route sezm_spin through SeZMSpinModel."""
-        model = get_model(self._build_model_params()).to(self.device)
-        data = {
-            "model": model.serialize(),
-            "model_def_script": self._build_model_params(),
-            "@variables": {},
-        }
-
-        with (
-            tempfile.TemporaryDirectory() as tmpdir,
-            mock.patch(
-                "deepmd.pt.utils.serialization.torch.jit.script",
-                side_effect=lambda model: model,
-            ),
-            mock.patch("deepmd.pt.utils.serialization.torch.jit.save") as save_mock,
-        ):
-            deserialize_to_file(f"{tmpdir}/model.pth", data)
-
-        saved_model = save_mock.call_args.args[0]
-        self.assertIsInstance(saved_model, SeZMSpinModel)
-        self.assertEqual(
-            saved_model.model_def_script,
-            json.dumps(data["model_def_script"]),
-        )
-
-    def test_energy_spin_loss_consumes_force_mag(self) -> None:
-        """EnergySpinLoss should consume force and magnetic-force predictions."""
-        model = get_model(self._build_model_params()).to(self.device)
-        loss = EnergySpinLoss(
-            start_pref_e=1.0,
-            limit_pref_e=1.0,
-            start_pref_fr=1.0,
-            limit_pref_fr=1.0,
-            start_pref_fm=1.0,
-            limit_pref_fm=1.0,
-        )
-        input_dict = {
-            "coord": self.coord,
-            "atype": self.atype,
-            "spin": self.spin,
-            "box": self.box,
-        }
-        label = {
-            "energy": torch.zeros((1, 1), dtype=torch.float64, device=self.device),
-            "force": torch.zeros((1, 3, 3), dtype=torch.float64, device=self.device),
-            "force_mag": torch.zeros(
-                (1, 3, 3), dtype=torch.float64, device=self.device
-            ),
-            "find_energy": torch.tensor(1.0, device=self.device),
-            "find_force": torch.tensor(1.0, device=self.device),
-            "find_force_mag": torch.tensor(1.0, device=self.device),
-        }
-
-        model_pred, loss_value, more_loss = loss(
-            input_dict,
-            model,
-            label,
-            natoms=3,
-            learning_rate=1.0,
-        )
-
-        self.assertIn("force_mag", model_pred)
-        self.assertIn("rmse_fm", more_loss)
-        self.assertTrue(torch.isfinite(loss_value))
-
-    def test_dens_mode_is_rejected(self) -> None:
-        """SeZM spin permanently rejects the dens path."""
-        model = get_model(self._build_model_params()).to(self.device)
-
-        with self.assertRaises(NotImplementedError):
-            prepare_model_for_loss(model, {"type": "dens"})
-
-    def test_bridging_is_rejected(self) -> None:
-        """The virtual-atom scheme does not compose with an analytical potential.
-
-        A magnetic atom sits within the bridging window of its own virtual
-        partner, so a bridged descriptor would freeze every magnetic atom; the
-        model refuses the combination instead of building it.
-        """
-        with self.assertRaisesRegex(ValueError, "native"):
-            get_model(self._build_model_params(bridging_method="ZBL"))
-
-    @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
-    def test_compile_matches_eager(self) -> None:
-        """Compiled SeZM spin path should match eager predictions."""
-        eager = get_model(self._build_model_params(use_compile=False)).to(self.device)
-        with mock.patch.dict(os.environ, {"DP_COMPILE_INFER": "1"}, clear=False):
-            compiled = get_model(self._build_model_params(use_compile=True)).to(
-                self.device
-            )
-        compiled.load_state_dict(copy.deepcopy(eager.state_dict()))
-        eager.eval()
-        compiled.eval()
-
-        out_eager = eager(self.coord, self.atype, spin=self.spin, box=self.box)
-        out_compiled = compiled(self.coord, self.atype, spin=self.spin, box=self.box)
-
-        self.assertIn((False, True), compiled.compiled_core_compute_cache)
-        _assert_close_with_strict_warning(
-            out_compiled["energy"],
-            out_eager["energy"],
-            atol=1.0e-6,
-            rtol=1.0e-6,
-            msg="spin compile energy mismatch",
-        )
-        _assert_close_with_strict_warning(
-            out_compiled["force"],
-            out_eager["force"],
-            atol=1.0e-6,
-            rtol=1.0e-6,
-            msg="spin compile force mismatch",
-        )
-        _assert_close_with_strict_warning(
-            out_compiled["force_mag"],
-            out_eager["force_mag"],
-            atol=1.0e-6,
-            rtol=1.0e-6,
-            msg="spin compile magnetic force mismatch",
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/source/tests/pt/model/test_sezm_vacuum_freeze.py b/source/tests/pt/model/test_sezm_vacuum_freeze.py
index dc3a9766ec..bb1bce9490 100644
--- a/source/tests/pt/model/test_sezm_vacuum_freeze.py
+++ b/source/tests/pt/model/test_sezm_vacuum_freeze.py
@@ -27,12 +27,14 @@
     ModelWrapper,
 )
 
+from ...common.sezm_config import (
+    _tiny_sezm_model_params,
+)
 from .test_sezm_export import (
     _CPU,
     _SKIP_OFF_COMPILE_TORCH,
     _SKIP_OFF_COMPILE_TORCH_REASON,
     _clear_default_device,
-    _tiny_sezm_model_params,
 )
 
 BIAS = np.array([[-3.0], [0.5]])
diff --git a/source/tests/pt/model/test_sezm_vacuum_ref.py b/source/tests/pt/model/test_sezm_vacuum_ref.py
index 82e9781d46..a77d9d2d97 100644
--- a/source/tests/pt/model/test_sezm_vacuum_ref.py
+++ b/source/tests/pt/model/test_sezm_vacuum_ref.py
@@ -268,7 +268,7 @@ def test_fold_drops_the_compiled_graphs(self) -> None:
 
         model = self.make_model(True)
         key_live = _sezm_structure_key(model)
-        model.compiled_core_compute_cache[(False, False)] = object()
+        model.compiled_core_compute_cache[False] = object()
         model.fold_vacuum_reference()
         self.assertEqual(model.compiled_core_compute_cache, {})
         # a folded model traces a graph without reference nodes
diff --git a/source/tests/pt/model/test_spin_model_virtual_types.py b/source/tests/pt/model/test_spin_model_virtual_types.py
index b8d8dbf772..8193751b06 100644
--- a/source/tests/pt/model/test_spin_model_virtual_types.py
+++ b/source/tests/pt/model/test_spin_model_virtual_types.py
@@ -61,6 +61,19 @@ def _tensor(array: np.ndarray) -> torch.Tensor:
 class TestPtSpinModelVirtualTypes(unittest.TestCase):
     """Every pt lookup must mask ``atype < 0`` exactly as dpmodel does."""
 
+    def test_real_type_sequence_bias_override(self) -> None:
+        model = get_pt_model(MODEL_CONFIG)
+        sample = {
+            "coord": _tensor(np.arange(9, dtype=np.float64).reshape(1, 3, 3)),
+            "atype": _tensor(np.array([[0, 1, 2]], dtype=np.int64)),
+            "spin": _tensor(np.ones((1, 3, 3), dtype=np.float64)),
+        }
+        model.change_out_bias([sample], preset_out_bias={"energy": [1.0, 2.0, 3.0]})
+        np.testing.assert_allclose(
+            to_numpy_array(model.get_out_bias()).reshape(-1),
+            [1.0, 2.0, 3.0, 0.0, 0.0, 0.0],
+        )
+
     def test_dense_spin_expansion_preserves_virtual_types(self) -> None:
         model = get_pt_model(MODEL_CONFIG)
         coord = np.arange(9, dtype=np.float64).reshape(1, 3, 3)
diff --git a/source/tests/pt/test_auto_batch_size.py b/source/tests/pt/test_auto_batch_size.py
index 1927918440..0b970e3074 100644
--- a/source/tests/pt/test_auto_batch_size.py
+++ b/source/tests/pt/test_auto_batch_size.py
@@ -23,287 +23,276 @@
 )
 
 
-class _CudaAllocator:
-    """Deterministic allocator with retained workspace and graph-private memory."""
-
-    def __init__(self) -> None:
-        self.total = 1000
-        self.live = 100
-        self.private = 0
-        self.private_active = 0
-        self.retained = self.live
-        self.reserved = self.live
-        self.peak = self.live
-        self.allocated_total = self.live
-        self.allocations = 0
-        self.frees = 0
-        self.ooms = 0
-        self.frame_bytes = 100
-        self.batches = []
-
-    def stats(self, device: int) -> dict:
-        return {
-            "allocated_bytes": {
-                "all": {
-                    "current": self.live + self.private_active,
-                    "peak": self.peak,
-                    "allocated": self.allocated_total,
-                }
-            },
-            "active_bytes": {"all": {"current": self.live + self.private_active}},
-            "reserved_bytes": {"all": {"current": self.reserved + self.private}},
-            "reserved_bytes_by_private_pools": {
-                (0, 1): {"all": {"current": self.private}}
-            },
-            "num_device_alloc": self.allocations,
-            "num_device_free": self.frees,
-        }
-
-    def mem_get_info(self, device: int) -> tuple[int, int]:
-        return self.total - self.reserved - self.private, self.total
-
-    def empty_cache(self) -> None:
-        if self.reserved > self.retained:
-            self.reserved = self.retained
-            self.frees += 1
-
-    def snapshot(self) -> list[dict]:
-        return [
-            {
-                "device": 1,
-                "segment_pool_id": (0, 0),
-                "total_size": self.reserved,
-                "active_size": self.live,
-            },
-            {
-                "device": 1,
-                "segment_pool_id": (0, 1),
-                "total_size": self.private,
-                "active_size": self.private_active,
-            },
-            {
-                "device": 0,
-                "segment_pool_id": (0, 0),
-                "total_size": 10000,
-                "active_size": 0,
-            },
-        ]
-
-    def evaluate(self, data: np.ndarray) -> np.ndarray:
-        self.batches.append(len(data))
-        required = self.live + len(data) * self.frame_bytes
-        if required + self.private > self.total:
-            self.ooms += 1
-            raise torch.cuda.OutOfMemoryError("CUDA out of memory.")
-        if required > self.reserved:
-            self.reserved = required
-            self.allocations += 1
-        self.peak = max(self.peak, required + self.private_active)
-        self.allocated_total += len(data) * self.frame_bytes
-        return data.copy()
-
-
-class TestCudaMemoryBatching(unittest.TestCase):
+class TestCudaAdaptiveBatching(unittest.TestCase):
     def setUp(self) -> None:
-        self.allocator = _CudaAllocator()
         stack = ExitStack()
         self.addCleanup(stack.close)
         stack.enter_context(mock.patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "0"}))
-        stack.enter_context(
-            mock.patch(
-                "deepmd.pt.utils.auto_batch_size.env.DEVICE", torch.device("cuda:1")
-            )
-        )
         stack.enter_context(mock.patch("torch.cuda.is_available", return_value=True))
         stack.enter_context(
-            mock.patch("torch.cuda.get_allocator_backend", return_value="native")
-        )
-        stack.enter_context(mock.patch("torch.cuda.current_device", return_value=1))
-        self.device = stack.enter_context(
-            mock.patch("torch.cuda.device", return_value=nullcontext())
-        )
-        self.stats = stack.enter_context(
-            mock.patch(
-                "torch.cuda.memory_stats_as_nested_dict",
-                side_effect=self.allocator.stats,
-            )
+            mock.patch("deepmd.pt.utils.env.DEVICE", torch.device("cuda:0"))
         )
+        stack.enter_context(mock.patch("torch.cuda.device", return_value=nullcontext()))
+        stack.enter_context(mock.patch("torch.cuda.current_device", return_value=0))
         stack.enter_context(
-            mock.patch(
-                "torch.cuda.mem_get_info", side_effect=self.allocator.mem_get_info
-            )
-        )
-        stack.enter_context(
-            mock.patch("torch.cuda.get_per_process_memory_fraction", return_value=1.0)
-        )
-        self.empty_cache = stack.enter_context(
-            mock.patch("torch.cuda.empty_cache", side_effect=self.allocator.empty_cache)
+            mock.patch("torch.cuda.get_allocator_backend", return_value="native")
         )
-        stack.enter_context(
-            mock.patch(
-                "torch.cuda.memory_snapshot", side_effect=self.allocator.snapshot
+        # A successful hot path must not inspect or mutate allocator state.
+        for name in (
+            "memory_stats_as_nested_dict",
+            "memory_snapshot",
+            "mem_get_info",
+            "reset_peak_memory_stats",
+            "set_per_process_memory_fraction",
+        ):
+            stack.enter_context(
+                mock.patch.object(
+                    torch.cuda, name, side_effect=AssertionError(f"unexpected {name}")
+                )
             )
-        )
-        self.reset_peak = stack.enter_context(
-            mock.patch("torch.cuda.reset_peak_memory_stats")
-        )
-        self.set_fraction = stack.enter_context(
-            mock.patch("torch.cuda.set_per_process_memory_fraction")
-        )
+        self.empty_cache = stack.enter_context(mock.patch("torch.cuda.empty_cache"))
 
-    def _evaluate(self, auto: AutoBatchSize, natoms: int = 2) -> np.ndarray:
-        data = np.arange(128).reshape(64, 2)
-        result = auto.execute_all(self.allocator.evaluate, len(data), natoms, data)
-        np.testing.assert_array_equal(result, data)
-        return result
-
-    def test_growth_stops_before_allocator_oom(self) -> None:
-        """The same workload crosses OOM under unbounded doubling, but not admission."""
-        legacy = AutoBatchSize(128)
-        data = np.arange(128).reshape(64, 2)
-        with mock.patch.object(legacy, "_get_batch_budget", return_value=None):
-            legacy.execute_all(self.allocator.evaluate, len(data), 2, data)
-        self.assertGreater(self.allocator.ooms, 0)
-        self.allocator.empty_cache()
-        self.allocator.batches.clear()
-        self.allocator.ooms = 0
-        self._evaluate(AutoBatchSize(128))
-        self.assertEqual(self.allocator.ooms, 0)
-        self.assertEqual(self.allocator.batches[0], 1)
-        self.assertEqual(max(self.allocator.batches), 8)
-        self.device.assert_called_with(torch.device("cuda:1"))
-        self.reset_peak.assert_not_called()
-        self.set_fraction.assert_not_called()
-
-    def test_private_pool_is_not_available_workspace(self) -> None:
-        self.allocator.private = 500
-        for private_active in (0, 500):
-            with self.subTest(private_active=private_active):
-                self.allocator.private_active = private_active
-                self.allocator.batches.clear()
-                self._evaluate(AutoBatchSize(128))
-                self.assertEqual(self.allocator.ooms, 0)
-                self.assertEqual(max(self.allocator.batches), 3)
-
-    def test_calibration_preserves_the_requested_atom_budget(self) -> None:
+    def test_short_systems_reuse_atom_budget_without_calibration(self) -> None:
+        auto = AutoBatchSize(1024)
         data = np.arange(16).reshape(8, 2)
-        result = AutoBatchSize(16).execute_all(
-            self.allocator.evaluate, len(data), 2, data
-        )
-        np.testing.assert_array_equal(result, data)
-        self.assertEqual(self.allocator.batches, [1, 7])
-        self.assertEqual(self.allocator.ooms, 0)
-
-    def test_calibration_survives_a_stale_global_peak(self) -> None:
-        self.allocator.peak = self.allocator.total
-        self._evaluate(AutoBatchSize(128))
-        self.assertEqual(self.allocator.ooms, 0)
-        self.assertEqual(max(self.allocator.batches), 8)
-        self.assertEqual(self.allocator.peak, self.allocator.total)
-
-    def test_calibration_accounts_for_unreleasable_cache(self) -> None:
-        self.allocator.reserved = self.allocator.retained = 700
-        self.allocator.peak = self.allocator.total
-        self._evaluate(AutoBatchSize(128))
-        self.assertEqual(self.allocator.ooms, 0)
-        self.assertEqual(max(self.allocator.batches), 8)
-        self.assertEqual(self.allocator.peak, self.allocator.total)
-
-    def test_reuses_calibration_and_remeasures_changed_workload(self) -> None:
-        auto = AutoBatchSize(128)
-        self._evaluate(auto)
-        cache_releases = self.empty_cache.call_count
-        self.allocator.batches.clear()
-        self._evaluate(auto)
-        self.assertEqual(self.allocator.batches[0], 8)
-        self.assertEqual(self.empty_cache.call_count, cache_releases)
-        self.allocator.frame_bytes = 150
-        self.allocator.batches.clear()
-        self._evaluate(auto, natoms=4)
-        self.assertEqual(self.allocator.batches[0], 1)
-        self.assertEqual(max(self.allocator.batches), 5)
-        self.assertEqual(self.allocator.ooms, 0)
-
-    def test_recalibrates_after_training_claims_device_memory(self) -> None:
-        auto = AutoBatchSize(128)
-        self._evaluate(auto)
-        self.allocator.empty_cache()
-        self.allocator.private = 700
-        self.allocator.allocations += 1
-        self.allocator.batches.clear()
-        self._evaluate(auto)
-        self.assertEqual(max(self.allocator.batches), 1)
-        self.assertEqual(self.allocator.ooms, 0)
-
-    def test_fixed_budget_and_single_frames_do_not_profile(self) -> None:
-        with mock.patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "8"}):
-            self._evaluate(AutoBatchSize())
-        self.assertEqual(self.allocator.batches[0], 4)
-        self.stats.assert_not_called()
-        auto = AutoBatchSize()
-        data = np.ones((1, 2))
-        auto.execute_all(self.allocator.evaluate, 1, 2, data)
-        self.stats.assert_not_called()
+        attempts = []
 
-    def test_single_frame_oom_and_unrelated_errors_propagate(self) -> None:
-        self.allocator.frame_bytes = self.allocator.total
-        with self.assertRaises(OutOfMemoryError):
-            self._evaluate(AutoBatchSize())
-        failure = RuntimeError("invalid model input")
-        with self.assertRaisesRegex(RuntimeError, "invalid model input"):
-            AutoBatchSize().execute_all(
-                mock.Mock(side_effect=failure), 2, 2, np.ones((2, 2))
-            )
+        def evaluate(batch: np.ndarray) -> np.ndarray:
+            attempts.append(len(batch))
+            return batch.copy()
+
+        for natoms in (100, 101, 100, 101):
+            result = auto.execute_all(evaluate, len(data), natoms, data)
+            np.testing.assert_array_equal(result, data)
+        self.assertEqual(attempts, [8, 8, 8, 8])
+        self.assertEqual(auto.current_batch_size, 1024)
+        self.empty_cache.assert_not_called()
+
+    def test_fixed_workspace_does_not_trap_batch_at_one_frame(self) -> None:
+        auto = AutoBatchSize(100)
+        data = np.arange(128).reshape(64, 2)
+        attempts = []
 
-    def test_unexpected_workspace_increase_keeps_oom_backoff(self) -> None:
-        auto = AutoBatchSize(128)
+        def evaluate(batch: np.ndarray) -> np.ndarray:
+            attempts.append(len(batch))
+            # Six frames fit; one-frame extrapolation would admit only one.
+            if 600 + 50 * len(batch) > 900:
+                raise torch.cuda.OutOfMemoryError()
+            return batch.copy()
+
+        for _ in range(2):
+            result = auto.execute_all(evaluate, len(data), 100, data)
+            np.testing.assert_array_equal(result, data)
+        self.assertEqual(attempts[:5], [1, 2, 4, 8, 4])
+        self.assertIn(auto.current_batch_size, (400, 800))
+        self.assertEqual(auto._reprobe_interval, 8)
+        self.assertGreater(attempts.count(4), attempts.count(8))
+        self.assertEqual(self.empty_cache.call_count, attempts.count(8))
+
+    def test_changed_memory_pressure_backs_off_and_reuses_the_bound(self) -> None:
+        auto = AutoBatchSize(8)
         data = np.arange(128).reshape(64, 2)
+        capacity = 8
+        attempts = []
 
         def evaluate(batch: np.ndarray) -> np.ndarray:
-            if batch[0, 0] >= 2:
-                self.allocator.frame_bytes = 200
-            return self.allocator.evaluate(batch)
+            attempts.append(len(batch))
+            if len(batch) > capacity:
+                raise torch.cuda.OutOfMemoryError()
+            return batch.copy()
 
         result = auto.execute_all(evaluate, len(data), 2, data)
         np.testing.assert_array_equal(result, data)
-        self.assertEqual(self.allocator.ooms, 1)
+        self.assertIn(auto.current_batch_size, (16, 32))
+        capacity = 2
+        attempts.clear()
+        for _ in range(2):
+            result = auto.execute_all(evaluate, len(data), 2, data)
+            np.testing.assert_array_equal(result, data)
+        first_success = attempts.index(2)
+        self.assertGreaterEqual(attempts[0], 8)
+        self.assertEqual(attempts[first_success - 1], 4)
+        self.assertLessEqual(max(attempts[first_success:]), 4)
+        self.assertGreater(attempts.count(2), attempts.count(4))
+        self.assertEqual(auto._reprobe_interval, 8)
+
+    def test_dense_to_sparse_recovers_even_at_maximum_probe_interval(self) -> None:
+        auto = AutoBatchSize(2)
+        capacity = 2
+        attempts = []
+
+        def evaluate(nframes: int, start: int) -> tuple[int, None]:
+            attempts.append((nframes, start))
+            if nframes > capacity:
+                raise torch.cuda.OutOfMemoryError()
+            return nframes, None
+
+        # A stationary dense workload damps repeated probes, but not forever.
+        for _ in range(80):
+            auto.execute(evaluate, 0, 1)
+        self.assertEqual(auto._reprobe_interval, 8)
+        failures = [i for i, (n, _) in enumerate(attempts) if n > capacity]
+        self.assertTrue(
+            all(b - a == 9 for a, b in zip(failures[-4:-1], failures[-3:], strict=True))
+        )
+        capacity = 16  # Same atom count, less memory per frame.
+        recovered = False
+        for _ in range(11):  # At most eight successes, then 4 -> 8 -> 16.
+            n, _ = auto.execute(evaluate, 0, 1)
+            if n == capacity:
+                recovered = True
+                break
+        self.assertTrue(recovered)
+        self.assertEqual(auto._reprobe_interval, 8)
+
+    def test_initial_backoff_and_clipped_probe_preserve_the_failed_bound(self) -> None:
+        auto = AutoBatchSize(64)
+        capacity = 2
+        attempts = []
+
+        def evaluate(nframes: int, start: int) -> tuple[int, None]:
+            attempts.append((nframes, start))
+            if nframes > capacity:
+                raise torch.cuda.OutOfMemoryError()
+            return nframes, None
+
+        for _ in range(5):
+            self.assertEqual(auto.execute(evaluate, 17, 1), (0, None))
+        self.assertEqual(attempts, [(n, 17) for n in (64, 32, 16, 8, 4)])
+        self.assertEqual(auto._reprobe_interval, 2)
+        auto.execute(evaluate, 17, 1)
+        auto.execute(evaluate, 19, 1)
+        self.assertEqual(auto.current_batch_size, 4)
+        auto.execute(evaluate, 21, 1, max_nframes=1)
+        self.assertEqual(auto.minimal_not_working_batch_size, 4)
+        self.assertEqual(auto.current_batch_size, 4)
+        capacity = 8
+        auto.execute(evaluate, 22, 1)
+        self.assertEqual(auto.current_batch_size, 8)
+        self.assertGreater(auto.minimal_not_working_batch_size, 8)
+
+    def test_short_tails_and_atom_count_changes_do_not_age_a_probe(self) -> None:
+        auto = AutoBatchSize(400)
+        auto.execute(mock.Mock(side_effect=torch.cuda.OutOfMemoryError()), 0, 100)
+        self.assertEqual(auto.current_batch_size, 200)
+        for natoms in (50, 51) * 8:
+            auto.execute(lambda n, start: (n, None), 0, natoms, max_nframes=1)
+        self.assertEqual(auto.current_batch_size, 200)
+        self.assertEqual(auto.minimal_not_working_batch_size, 400)
+        self.assertEqual(auto._successful_batches, 0)
+        self.assertEqual(auto._reprobe_interval, 2)
+
+    def test_alternating_capacity_does_not_reset_probe_interval(self) -> None:
+        auto = AutoBatchSize(2)
+        capacity = 2
+
+        def evaluate(nframes: int, start: int) -> tuple[int, None]:
+            if nframes > capacity:
+                raise torch.cuda.OutOfMemoryError()
+            return nframes, None
+
+        auto.execute(evaluate, 0, 1)  # 2 -> 4.
+        auto.execute(evaluate, 0, 1)  # First dense OOM.
+        for interval, next_interval in ((2, 4), (4, 8), (8, 8)):
+            for _ in range(interval):
+                auto.execute(evaluate, 0, 1)
+            self.assertEqual(auto.current_batch_size, 4)
+            self.assertEqual(auto._reprobe_interval, next_interval)
+            capacity = 4
+            auto.execute(evaluate, 0, 1)  # Sparse probe succeeds and grows to 8.
+            self.assertEqual(auto.current_batch_size, 8)
+            capacity = 2
+            auto.execute(evaluate, 0, 1)
+            auto.execute(evaluate, 0, 1)  # Dense backoff: 8 -> 4 -> 2.
+            self.assertEqual(auto.current_batch_size, 2)
+            self.assertEqual(auto._reprobe_interval, next_interval)
+
+    def test_automatic_retry_mode_completes_before_reprobing(self) -> None:
+        auto = AutoBatchSize(4)
+        data = np.arange(82).reshape(41, 2)
+        capacity = 2
+        attempts = []
 
-    def test_respects_existing_allocator_memory_fraction(self) -> None:
-        with mock.patch("torch.cuda.get_per_process_memory_fraction", return_value=0.5):
-            self._evaluate(AutoBatchSize(128))
-        self.assertEqual(max(self.allocator.batches), 3)
-        self.set_fraction.assert_not_called()
+        def evaluate(batch: np.ndarray) -> np.ndarray:
+            attempts.append((len(batch), int(batch[0, 0])))
+            if len(batch) > capacity:
+                raise torch.cuda.OutOfMemoryError()
+            return batch.copy()
+
+        def evaluate_with_restart() -> tuple[np.ndarray, int]:
+            # Match hook-based feature extraction: retry the entire input and
+            # restore the mode in finally, including on unsuccessful attempts.
+            for attempt in range(3):
+                auto.set_oom_retry_mode(True)
+                try:
+                    return auto.execute_all(evaluate, len(data), 1, data), attempt
+                except RetrySignal:
+                    continue
+                finally:
+                    auto.set_oom_retry_mode(False)
+            self.fail("whole-input retry did not make bounded progress")
+
+        result, retries = evaluate_with_restart()
+        np.testing.assert_array_equal(result, data)
+        self.assertEqual(retries, 1)
+        self.assertEqual(attempts[0], (4, 0))
+        self.assertEqual(attempts[1], (2, 0))
+        # Schedule the next invocation's probe only after this pass completes,
+        # including when its last batch is a short tail.
+        self.assertEqual(auto.current_batch_size, 4)
+        self.assertEqual(auto.minimal_not_working_batch_size, 4)
+        for _ in range(3):
+            result, retries = evaluate_with_restart()
+            np.testing.assert_array_equal(result, data)
+            self.assertEqual(retries, 1)
+        capacity = 8
+        attempts.clear()
+        result, retries = evaluate_with_restart()
+        np.testing.assert_array_equal(result, data)
+        self.assertEqual(retries, 1)  # 16 fails once; restarted 8-frame pass fits.
+        self.assertTrue(any(n == 8 for n, _ in attempts))
 
-    def test_retry_signal_preserves_backoff(self) -> None:
-        auto = AutoBatchSize(8)
+    def test_retry_mode_does_not_schedule_probe_when_concatenation_fails(self) -> None:
+        auto = AutoBatchSize(4)
+        auto.execute(mock.Mock(side_effect=torch.cuda.OutOfMemoryError()), 0, 1)
         auto.set_oom_retry_mode(True)
-        data = np.arange(128).reshape(64, 2)
+        with self.assertRaisesRegex(RuntimeError, "Unexpected result type"):
+            auto.execute_all(lambda batch: object(), 41, 1, np.ones((41, 2)))
+        self.assertEqual(auto.current_batch_size, 2)
+        self.assertEqual(auto.minimal_not_working_batch_size, 4)
+        self.assertEqual(auto._reprobe_interval, 2)
+
+    def test_explicit_limit_keeps_oom_recovery_and_retry_mode(self) -> None:
+        with mock.patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "8"}):
+            auto = AutoBatchSize()
+        auto.set_oom_retry_mode(True)
+        data = np.arange(16).reshape(8, 2)
+        attempts = []
 
         def evaluate(batch: np.ndarray) -> np.ndarray:
+            attempts.append(len(batch))
             if len(batch) > 2:
-                self.allocator.frame_bytes = 350
-            return self.allocator.evaluate(batch)
+                raise torch.cuda.OutOfMemoryError()
+            return batch.copy()
 
         with self.assertRaises(RetrySignal):
             auto.execute_all(evaluate, len(data), 2, data)
         self.assertEqual(auto.current_batch_size, 4)
-        result = auto.execute_all(self.allocator.evaluate, len(data), 2, data)
+        result = auto.execute_all(evaluate, len(data), 2, data)
         np.testing.assert_array_equal(result, data)
-        self.assertEqual(self.allocator.ooms, 1)
-
-    def test_cpu_and_async_allocator_keep_the_existing_policy(self) -> None:
-        data = np.ones((8, 2))
-        for device, backend in (("cpu", "native"), ("cuda:1", "cudaMallocAsync")):
-            with (
-                self.subTest(device=device, backend=backend),
-                mock.patch(
-                    "deepmd.pt.utils.auto_batch_size.env.DEVICE", torch.device(device)
-                ),
-                mock.patch("torch.cuda.get_allocator_backend", return_value=backend),
-            ):
-                result = AutoBatchSize().execute_all(np.copy, len(data), 2, data)
-                np.testing.assert_array_equal(result, data)
-        self.stats.assert_not_called()
+        self.assertEqual(attempts, [4, 2, 2, 2, 2])
+        self.empty_cache.assert_called_once()
+
+    def test_single_frame_oom_and_unrelated_errors_propagate(self) -> None:
+        data = np.ones((1, 2))
+        with self.assertRaises(OutOfMemoryError):
+            AutoBatchSize().execute_all(
+                mock.Mock(side_effect=torch.cuda.OutOfMemoryError()), 1, 2, data
+            )
+        failure = RuntimeError("invalid model input")
+        with self.assertRaisesRegex(RuntimeError, "invalid model input"):
+            AutoBatchSize().execute_all(mock.Mock(side_effect=failure), 1, 2, data)
+        self.empty_cache.assert_called_once()
 
 
 class TestAutoBatchSize(unittest.TestCase):
diff --git a/source/tests/pt/test_change_bias.py b/source/tests/pt/test_change_bias.py
index 746656ac8c..f7e3060293 100644
--- a/source/tests/pt/test_change_bias.py
+++ b/source/tests/pt/test_change_bias.py
@@ -144,6 +144,32 @@ def test_change_bias_with_data_sys_file(self) -> None:
         expected_bias = expected_model.get_out_bias()
         torch.testing.assert_close(updated_bias, expected_bias)
 
+    def test_change_bias_ignores_checkpoint_preset(self) -> None:
+        checkpoint = torch.load(
+            str(self.model_path), map_location=DEVICE, weights_only=True
+        )
+        preset = {"energy": [1.0e6] * len(self.config["model"]["type_map"])}
+        checkpoint["model"]["_extra_state"]["model_params"]["preset_out_bias"] = preset
+        torch.save(checkpoint, self.model_path)
+        run_dp(
+            f"dp --pt change-bias {self.model_path!s} -s {self.data_file[0]} -o {self.model_path_data_bias!s}"
+        )
+        updated = torch.load(
+            str(self.model_path_data_bias), map_location=DEVICE, weights_only=True
+        )
+        model_params = updated["model"]["_extra_state"]["model_params"]
+        self.assertEqual(model_params["preset_out_bias"], preset)
+        wrapper = ModelWrapper(get_model_for_wrapper(model_params))
+        wrapper.load_state_dict(updated["model"])
+        expected_model = model_change_out_bias(
+            self.trainer.wrapper.model["Default"],
+            self.sampled,
+            _bias_adjust_mode="change-by-statistic",
+        )
+        torch.testing.assert_close(
+            wrapper.model["Default"].get_out_bias(), expected_model.get_out_bias()
+        )
+
     def test_change_bias_with_user_defined(self) -> None:
         user_bias = [0.1, 3.2, -0.5]
         run_dp(
diff --git a/source/tests/pt/test_compile_compat.py b/source/tests/pt/test_compile_compat.py
index 12772edcd7..f33796c739 100644
--- a/source/tests/pt/test_compile_compat.py
+++ b/source/tests/pt/test_compile_compat.py
@@ -43,6 +43,7 @@ def test_fusion_size_defaults_to_eight(monkeypatch) -> None:
     assert inference_options["max_fusion_size"] == 8
     assert inference_options["cpp.min_chunk_size"] == 1
     assert inference_options["cpp.dynamic_threads"] is True
+    assert "cpp.max_horizontal_fusion_size" not in inference_options
 
 
 def test_torch_211_inference_keeps_default_cpp_parallelism(monkeypatch) -> None:
@@ -53,6 +54,49 @@ def test_torch_211_inference_keeps_default_cpp_parallelism(monkeypatch) -> None:
 
     assert "cpp.min_chunk_size" not in inference_options
     assert "cpp.dynamic_threads" not in inference_options
+    assert inference_options["cpp.max_horizontal_fusion_size"] == 1
+    assert "cpp.simdlen" not in inference_options
+    assert "cpp.max_horizontal_fusion_size" not in build_inductor_compile_options()
+
+
+def test_inductor_kernels_write_separate_output_buffers() -> None:
+    """Generated kernels write separate output buffers instead of an input."""
+
+    def fn(
+        a: torch.Tensor, b: torch.Tensor, w: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # The product is an extern kernel. Where in-place buffers are allowed,
+        # the fused kernel of the addition and the weighted sum stores the
+        # shifted values over the buffer of the product.
+        shifted = a @ b + 1.0
+        return shifted, (shifted * w).sum(-1)
+
+    from torch._inductor.utils import (
+        run_and_get_code,
+    )
+
+    a, b, w = (
+        torch.randn(shape, dtype=torch.float64, device="cpu")
+        for shape in ((16, 9), (9, 9), (16, 9))
+    )
+    expected = a @ b + 1.0
+
+    def generated_code(options: dict) -> list[str]:
+        compiled = torch.compile(fn, backend="inductor", options=options)
+        (shifted, total), codes = run_and_get_code(compiled, a, b, w)
+        torch.testing.assert_close(shifted, expected)
+        torch.testing.assert_close(total, (expected * w).sum(-1))
+        return codes
+
+    options = build_inductor_compile_options()
+    # The same graph with in-place buffers allowed must store over its input,
+    # so that the assertion below tests the option rather than the graph.
+    in_place = generated_code({**options, "inplace_buffers": True})
+    separate = generated_code(options)
+    torch._dynamo.reset()
+
+    assert any("in_out_ptr" in code for code in in_place)
+    assert not any("in_out_ptr" in code for code in separate)
 
 
 def test_fusion_size_environment_is_shared(monkeypatch) -> None:
diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py
index 4826048bce..36e31f7cd5 100644
--- a/source/tests/pt/test_dp_test.py
+++ b/source/tests/pt/test_dp_test.py
@@ -222,6 +222,45 @@ def test_dp_test_input_json_rglob(self) -> None:
         os.unlink(tmp_model_path)
         self.assertTrue(os.path.exists(self.detail_file + ".e.out"))
 
+    def test_training_groups_use_input_relative_patterns(self) -> None:
+        trainer = get_trainer(deepcopy(self.config))
+        model = torch.jit.script(trainer.model)
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            model_path = root / "model.pth"
+            torch.jit.save(model, str(model_path))
+            data_root = Path(__file__).parent / "water" / "data"
+            relative_root = os.path.relpath(data_root, root)
+            config = deepcopy(self.config)
+            config["training"]["training_data"] = [
+                {
+                    "systems": relative_root,
+                    "rglob_patterns": ["single"],
+                    "batch_size": "filter:1",
+                    "prob": 0.0,
+                },
+                {
+                    "systems": [relative_root],
+                    "rglob_patterns": ["data_0"],
+                    "prob": 1.0,
+                },
+            ]
+            input_path = root / "input.json"
+            input_path.write_text(json.dumps(config))
+            detail = str(root / "details")
+            dp_test(
+                model=str(model_path),
+                system=None,
+                datafile=None,
+                train_json=str(input_path),
+                numb_test=1,
+                rand_seed=None,
+                shuffle_test=False,
+                detail_file=detail,
+                atomic=False,
+            )
+            self.assertEqual(np.loadtxt(detail + ".e.out", ndmin=2).shape[0], 2)
+
     def tearDown(self) -> None:
         for f in os.listdir("."):
             if f.startswith("model") and f.endswith(".pt"):
diff --git a/source/tests/pt/test_dp_test_density.py b/source/tests/pt/test_dp_test_density.py
new file mode 100644
index 0000000000..65846427b7
--- /dev/null
+++ b/source/tests/pt/test_dp_test_density.py
@@ -0,0 +1,1265 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import json
+import os
+import shutil
+import tempfile
+import unittest
+from copy import (
+    deepcopy,
+)
+from pathlib import (
+    Path,
+)
+from unittest.mock import (
+    patch,
+)
+
+import numpy as np
+import torch
+
+from deepmd.entrypoints.test import test as dp_test
+from deepmd.infer.deep_density import (
+    DeepDensity,
+)
+from deepmd.infer.deep_eval import (
+    DeepEval,
+)
+from deepmd.infer.model_test import (
+    build_tester,
+)
+from deepmd.pt.entrypoints.main import (
+    get_trainer,
+)
+from deepmd.pt.model.atomic_model.density_atomic_model import (
+    DPDensityAtomicModel,
+)
+from deepmd.pt.model.atomic_model.dp_atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.pt.model.descriptor.hybrid import (
+    DescrptHybrid,
+)
+from deepmd.pt.model.descriptor.se_a import (
+    DescrptSeA,
+)
+from deepmd.pt.model.descriptor.se_r import (
+    DescrptSeR,
+)
+from deepmd.pt.model.model import (
+    get_model,
+)
+from deepmd.pt.model.model.density_model import (
+    GridDensityModel,
+)
+from deepmd.pt.model.task.density import (
+    DensityFittingNet,
+)
+from deepmd.utils.argcheck import (
+    normalize,
+)
+from deepmd.utils.data import (
+    DeepmdData,
+)
+from deepmd.utils.path import (
+    DPPath,
+)
+from deepmd.utils.stat_file import (
+    StatFileSpec,
+    open_stat_file,
+)
+
+model_density = {
+    # the last entry is the reserved grid point type; its sel is 0 because
+    # grid points are never neighbors, only centers
+    "type_map": ["O", "H", "X"],
+    "descriptor": {
+        "type": "se_e2_a",
+        "sel": [8, 8, 0],
+        "rcut_smth": 0.50,
+        "rcut": 4.00,
+        "neuron": [8, 16],
+        "resnet_dt": False,
+        "axis_neuron": 4,
+        # grid points may coincide with atoms; without protection the
+        # env matrix would divide by zero and produce NaN densities
+        "env_protection": 1e-6,
+        "seed": 1,
+    },
+    "fitting_net": {
+        "type": "density",
+        "neuron": [8, 8],
+        "seed": 1,
+    },
+}
+
+
+class TestDPTestDensity(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.tmpdir = tempfile.TemporaryDirectory()
+        root = Path(cls.tmpdir.name)
+
+        # write a synthetic density system
+        rng = np.random.default_rng(42)
+        cls.nframes, natoms, cls.ngrid = 4, 5, 8
+        cls.system = root / "system"
+        set_dir = cls.system / "set.000"
+        set_dir.mkdir(parents=True)
+        np.save(set_dir / "coord.npy", rng.random((cls.nframes, natoms * 3)) * 8.0)
+        np.save(
+            set_dir / "box.npy",
+            np.tile(np.eye(3).reshape(-1) * 10.0, (cls.nframes, 1)),
+        )
+        np.save(set_dir / "grid.npy", rng.random((cls.nframes, cls.ngrid, 3)) * 8.0)
+        np.save(set_dir / "density.npy", rng.random((cls.nframes, cls.ngrid, 1)))
+        np.savetxt(cls.system / "type.raw", [0, 0, 1, 1, 1], fmt="%d")
+        np.savetxt(cls.system / "type_map.raw", ["O", "H", "X"], fmt="%s")
+
+        cls.config = {
+            "model": deepcopy(model_density),
+            "learning_rate": {
+                "type": "exp",
+                "start_lr": 0.001,
+                "stop_lr": 1e-8,
+                "decay_steps": 10,
+            },
+            "optimizer": {"type": "Adam"},
+            "loss": {
+                "type": "grid_density",
+                "start_pref_d": 1.0,
+                "limit_pref_d": 1.0,
+            },
+            "training": {
+                "training_data": {
+                    "systems": [str(cls.system)],
+                    "batch_size": 1,
+                },
+                "validation_data": {
+                    "systems": [str(cls.system)],
+                    "batch_size": 1,
+                },
+                "numb_steps": 1,
+                "seed": 1,
+                "disp_file": os.devnull,
+                "save_freq": 100,
+            },
+        }
+
+        # build and freeze a tiny density model through the trainer path
+        trainer = get_trainer(normalize(deepcopy(cls.config)))
+        with torch.device("cpu"):
+            input_dict, _, _ = trainer.get_data(is_train=False)
+        # the density model takes grid instead of spin as the extra input
+        input_dict.pop("spin", None)
+        trainer.model(**input_dict)
+        # keep the live torch model and a data sample for the gradient check
+        cls.torch_model = trainer.model
+        cls.input_dict = input_dict
+        model = torch.jit.script(trainer.model)
+        tmp_fd, cls.model_path = tempfile.mkstemp(suffix=".pth")
+        os.close(tmp_fd)
+        torch.jit.save(model, cls.model_path)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        os.unlink(cls.model_path)
+        cls.tmpdir.cleanup()
+
+    def _stat_sample(self, module: torch.nn.Module) -> dict:
+        """Return a cloned sample on the device used by statistic collection."""
+        device = next(module.parameters()).device
+        return {
+            key: value.detach().clone().to(device)
+            if isinstance(value, torch.Tensor)
+            else deepcopy(value)
+            for key, value in self.input_dict.items()
+        }
+
+    def test_change_out_bias_noop(self) -> None:
+        # fine-tuning calls the model-level change_out_bias, whose default
+        # change-by-statistic mode would run the grid-less stat wrapper;
+        # for density models it must be a no-op instead
+        model = self.torch_model
+        bias_before = model.atomic_model.out_bias.detach().clone()
+        model.change_out_bias(self.input_dict)
+        model.change_out_bias(self.input_dict, bias_adjust_mode="set-by-statistic")
+        # Keep both no-op overrides compatible with their base APIs.
+        model.change_out_bias(self.input_dict, "set-by-statistic", {})
+        model.atomic_model.change_out_bias(
+            self.input_dict, None, "set-by-statistic", None, None
+        )
+        torch.testing.assert_close(model.atomic_model.out_bias, bias_before)
+
+    def test_pair_excl_involving_grid_type_applies_in_forward(self) -> None:
+        # exclusions involving the grid type must drop the neighbor in the
+        # directional grid-to-atom list too, not just in the atom nlist
+        model = self.torch_model
+        grid_type = len(self.config["model"]["type_map"]) - 1
+        input_dict = {
+            kk: vv.clone() if isinstance(vv, torch.Tensor) else vv
+            for kk, vv in self.input_dict.items()
+        }
+        model.atomic_model.reinit_pair_exclude([(0, grid_type)])
+        try:
+            out_excl = model(**input_dict)["density"]
+        finally:
+            model.atomic_model.reinit_pair_exclude([])
+        out_none = model(**input_dict)["density"]
+        self.assertFalse(torch.allclose(out_excl, out_none))
+
+    def test_atom_excl_does_not_mask_grid(self) -> None:
+        # grid points carry the reserved type X; excluding a real atom type
+        # must not zero the density predictions (it would if grid_type were 0)
+        model = self.torch_model
+        model.atomic_model.reinit_atom_exclude([0])
+        try:
+            input_dict = {
+                kk: vv.clone() if isinstance(vv, torch.Tensor) else vv
+                for kk, vv in self.input_dict.items()
+            }
+            out = model(**input_dict)
+            self.assertFalse(bool((out["density"] == 0).all()))
+        finally:
+            model.atomic_model.reinit_atom_exclude([])
+
+    def test_property_named_density_dispatches_to_property(self) -> None:
+        # the dispatch keys on the grid capability, not the output name, so
+        # a property fitting with property_name "density" is not hijacked
+        from deepmd.infer.deep_property import (
+            DeepProperty,
+        )
+        from deepmd.pt.model.model.property_model import (
+            PropertyModel,
+        )
+        from deepmd.pt.model.task.property import (
+            PropertyFittingNet,
+        )
+
+        descriptor = DescrptSeA(rcut=4.0, rcut_smth=0.5, sel=[8, 8])
+        fitting = PropertyFittingNet(
+            descriptor.get_ntypes(),
+            descriptor.get_dim_out(),
+            "density",
+            neuron=[8, 8],
+        )
+        model = PropertyModel(descriptor, fitting, type_map=["O", "H"])
+        fd, path = tempfile.mkstemp(suffix=".pth")
+        os.close(fd)
+        try:
+            with torch.device("cpu"):
+                torch.jit.save(torch.jit.script(model), path)
+            dp = DeepEval(path)
+            self.assertIsInstance(dp, DeepProperty)
+            # the property must stay evaluable: the eval guard keys on the
+            # grid capability, not the output-variable name
+            set_dir = self.system / "set.000"
+            coord = np.load(set_dir / "coord.npy")[:2]
+            box = np.load(set_dir / "box.npy")[:2]
+            atype = np.loadtxt(self.system / "type.raw", dtype=int)
+            result = dp.eval(coord, box, atype)
+            self.assertEqual(result[0].shape[0], 2)
+        finally:
+            os.unlink(path)
+
+    def _run_training(self, config: dict, workdir: str) -> None:
+        """Run a few real training steps (checkpoints land in workdir)."""
+        cwd = os.getcwd()
+        os.chdir(workdir)
+        try:
+            trainer = get_trainer(normalize(deepcopy(config)))
+            trainer.run()
+        finally:
+            os.chdir(cwd)
+        self.assertTrue(os.path.exists(os.path.join(workdir, "model.ckpt.pt")))
+
+    def _training_config(
+        self, system: Path, start: float = 1.0, limit: float = 1.0
+    ) -> dict:
+        config = deepcopy(self.config)
+        config["training"]["training_data"]["systems"] = [str(system)]
+        config["training"]["validation_data"]["systems"] = [str(system)]
+        config["training"]["training_data"]["batch_size"] = 2
+        config["training"]["validation_data"]["batch_size"] = 2
+        config["training"]["numb_steps"] = 3
+        config["loss"]["start_pref_d"] = start
+        config["loss"]["limit_pref_d"] = limit
+        return config
+
+    def test_eval_single_frame_grid_2d(self) -> None:
+        # a natural single-frame grid of shape (ngrid, 3) must be carried
+        # through with its frame dimension, not sliced to one grid point
+        dp = DeepDensity(self.model_path)
+        set_dir = self.system / "set.000"
+        coord = np.load(set_dir / "coord.npy")[:1]
+        box = np.load(set_dir / "box.npy")[:1]
+        grid = np.load(set_dir / "grid.npy")[0]  # (ngrid, 3): no frame dim
+        atype = np.loadtxt(self.system / "type.raw", dtype=int)
+        out = dp.eval(coord, box, atype, grid=grid)
+        self.assertEqual(out.shape, (1, self.ngrid))
+
+    def test_eval_rejects_invalid_grid_inputs(self) -> None:
+        dp = DeepDensity(self.model_path)
+        set_dir = self.system / "set.000"
+        coord = np.load(set_dir / "coord.npy")[:2]
+        box = np.load(set_dir / "box.npy")[:2]
+        grid = np.load(set_dir / "grid.npy")[:2]
+        atype = np.loadtxt(self.system / "type.raw", dtype=int)
+        invalid = (
+            ("frame mismatch", grid[:1], {}, "frames"),
+            ("ambiguous rank two", grid[0], {}, "ambiguous"),
+            ("wrong coordinate width", np.zeros((2, self.ngrid, 2)), {}, "shape"),
+            ("rank four", np.zeros((2, self.ngrid, 3, 1)), {}, "shape"),
+            ("empty", np.zeros((2, 0, 3)), {}, "at least one point"),
+            ("invalid batch size", grid, {"grid_batch_size": 0}, "grid_batch_size"),
+        )
+        for case, invalid_grid, kwargs, message in invalid:
+            with self.subTest(case=case):
+                with self.assertRaisesRegex(ValueError, message):
+                    dp.eval(coord, box, atype, grid=invalid_grid, **kwargs)
+
+    def test_fitting_rejects_aparam(self) -> None:
+        # grid descriptor rows have no per-atom parameters
+        descriptor = DescrptSeA(rcut=4.0, rcut_smth=0.5, sel=[8, 8, 0])
+        with self.assertRaisesRegex(ValueError, "aparam"):
+            DensityFittingNet(
+                descriptor.get_ntypes(),
+                descriptor.get_dim_out(),
+                numb_aparam=2,
+            )
+
+    def test_fitting_model_serialization_roundtrip(self) -> None:
+        # GeneralFitting.serialize() emits @version 5; the density fitting
+        # and the whole model must round-trip through serialize/deserialize,
+        # which is what model freezing relies on. The model must go through
+        # BaseModel.deserialize, which dispatches on the serialized fitting
+        # type ("density") and is what dp convert-backend uses.
+        from deepmd.pt.loss.charge import (
+            GridDensityLoss,
+        )
+        from deepmd.pt.model.model.model import (
+            BaseModel,
+        )
+
+        fitting = self.torch_model.atomic_model.fitting_net
+        restored_fit = DensityFittingNet.deserialize(fitting.serialize())
+        self._assert_state_dict_equal(fitting, restored_fit)
+        model = self.torch_model
+        restored_model = BaseModel.deserialize(model.serialize())
+        self.assertIsInstance(restored_model, GridDensityModel)
+        self._assert_state_dict_equal(model, restored_model)
+
+        loss_fn = GridDensityLoss(
+            starter_learning_rate=0.001, start_pref_d=2.0, limit_pref_d=0.5
+        )
+        data = loss_fn.serialize()
+        self.assertEqual(data["@class"], "GridDensityLoss")
+        restored_loss = GridDensityLoss.deserialize(data)
+        self.assertEqual(restored_loss.serialize(), data)
+
+    def _assert_state_dict_equal(
+        self, module0: torch.nn.Module, module1: torch.nn.Module
+    ) -> None:
+        sd0 = module0.state_dict()
+        sd1 = module1.state_dict()
+        self.assertEqual(set(sd0), set(sd1))
+        for key in sd0:
+            # Serialization restores tensors on the default device (CPU),
+            # while the source module may live on CUDA. This round-trip
+            # checks serialized values and metadata, not device placement.
+            torch.testing.assert_close(sd0[key].cpu(), sd1[key].cpu())
+
+    def test_loss_inference_mode(self) -> None:
+        # inference=True reports the metrics even with zero prefactors
+        from deepmd.pt.loss.charge import (
+            GridDensityLoss,
+        )
+        from deepmd.pt.utils import (
+            env,
+        )
+
+        loss_fn = GridDensityLoss(inference=True)
+
+        class FakeModel(torch.nn.Module):
+            def forward(self, **kwargs):
+                return {
+                    "density": torch.tensor([[[2.0], [4.0]]], device=env.DEVICE),
+                    "mask": torch.tensor(
+                        [[1, 1]], dtype=torch.int32, device=env.DEVICE
+                    ),
+                }
+
+        label = {
+            "density": torch.tensor([[[1.0], [2.0]]], device=env.DEVICE),
+            "find_density": 1.0,
+        }
+        _, loss, more_loss = loss_fn({}, FakeModel(), label, 1, 1.0)
+        self.assertIn("rmse_d", more_loss)
+        self.assertIn("mae_d", more_loss)
+        self.assertAlmostEqual(more_loss["rmse_d"].item(), ((1.0 + 4.0) / 2) ** 0.5)
+
+    def test_loss_masks_excluded_grid_points(self) -> None:
+        # excluded grid points (mask == 0) must not contribute to the
+        # residual, which is normalised by the per-frame mask sum
+        from deepmd.pt.loss.charge import (
+            GridDensityLoss,
+        )
+        from deepmd.pt.utils import (
+            env,
+        )
+
+        loss_fn = GridDensityLoss(start_pref_d=1.0, limit_pref_d=1.0)
+
+        class FakeModel(torch.nn.Module):
+            def forward(self, **kwargs):
+                return {
+                    "density": torch.tensor(
+                        [[[1.0], [2.0], [3.0], [4.0]]], device=env.DEVICE
+                    ),
+                    "mask": torch.tensor(
+                        [[1, 1, 0, 0]], dtype=torch.int32, device=env.DEVICE
+                    ),
+                }
+
+        label = {
+            "density": torch.tensor(
+                [[[1.0], [1.0], [100.0], [100.0]]], device=env.DEVICE
+            ),
+            "find_density": 1.0,
+        }
+        _, loss, more_loss = loss_fn({}, FakeModel(), label, 1, 1.0)
+        # only the two unmasked points count: (0^2 + 1^2) / 2
+        self.assertAlmostEqual(loss.item(), 0.5)
+        self.assertAlmostEqual(more_loss["rmse_d"].item(), 0.5**0.5)
+
+    def test_training_steps(self) -> None:
+        # exercises GridDensityLoss.forward on the main path
+        workdir = tempfile.mkdtemp(dir=self.tmpdir.name)
+        self._run_training(self._training_config(self.system), workdir)
+
+    def test_training_requires_density_label(self) -> None:
+        # density is the only supervision signal: a missing density.npy must
+        # abort training instead of silently optimising nothing
+        system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "system"
+        shutil.copytree(self.system, system)
+        (system / "set.000" / "density.npy").unlink()
+        with self.assertRaisesRegex(RuntimeError, "not found"):
+            get_trainer(normalize(self._training_config(system)))
+
+    def test_training_zero_prefactor(self) -> None:
+        # start_pref_d = limit_pref_d = 0 disables the density term;
+        # backward() must still work on the graph-connected zero loss
+        workdir = tempfile.mkdtemp(dir=self.tmpdir.name)
+        self._run_training(
+            self._training_config(self.system, start=0.0, limit=0.0), workdir
+        )
+
+    def test_training_with_disp_avg(self) -> None:
+        # TaskLoss.training_metric_names is consulted before training starts
+        # whenever averaged metrics are requested.
+        config = self._training_config(self.system)
+        config["training"]["disp_avg"] = True
+        self._run_training(config, tempfile.mkdtemp(dir=self.tmpdir.name))
+
+    def test_env_protection_default(self) -> None:
+        # a density model built without an explicit env_protection gets
+        # 1e-6 by default (with a warning), and the recorded def script
+        # agrees with the built model
+        config = deepcopy(self.config)
+        config["model"]["descriptor"].pop("env_protection")
+        trainer = get_trainer(normalize(config))
+        descriptor = trainer.model.atomic_model.descriptor
+        self.assertEqual(descriptor.get_env_protection(), 1e-6)
+        recorded = json.loads(trainer.model.model_def_script)
+        self.assertEqual(recorded["descriptor"]["env_protection"], 1e-6)
+
+    def test_hybrid_descriptor_stat_blocks(self) -> None:
+        # hybrid descriptors keep their blocks in descrpt_list: the block
+        # discovery must walk the module tree instead of a hard-coded
+        # attribute list, and a stat file must not abort the run
+        descriptor = DescrptHybrid(
+            [
+                DescrptSeA(rcut=4.0, rcut_smth=0.5, sel=[8, 8, 0]),
+                DescrptSeR(rcut=2.0, rcut_smth=1.0, sel=[4, 4, 0]),
+            ]
+        )
+        fitting = DensityFittingNet(descriptor.get_ntypes(), descriptor.get_dim_out())
+        model = DPDensityAtomicModel(descriptor, fitting, type_map=["O", "H", "X"])
+        blocks = model._descriptor_stat_blocks()
+        self.assertEqual(len(blocks), 2)
+        # a stat file must take the soft-failure path, not abort
+        stat_dir = DPPath(tempfile.mkdtemp(dir=self.tmpdir.name), "w")
+
+        def sampler() -> list:
+            return [self._stat_sample(model)]
+
+        model.compute_or_load_stat(sampler, stat_file_path=stat_dir)
+
+    def test_grid_type_requires_reserved_slot(self) -> None:
+        # a real element in the reserved slot must be rejected, not
+        # silently used as the grid type
+        descriptor = DescrptSeA(rcut=4.0, rcut_smth=0.5, sel=[8, 8])
+        fitting = DensityFittingNet(descriptor.get_ntypes(), descriptor.get_dim_out())
+        with self.assertRaisesRegex(ValueError, "reserved grid type"):
+            DPDensityAtomicModel(descriptor, fitting, type_map=["O", "H"])
+        # non-element reserved slot is accepted (all other tests use X)
+        DPDensityAtomicModel(descriptor, fitting, type_map=["O", "X"])
+
+    def test_env_protection_enforced_at_model_level(self) -> None:
+        # models constructed directly with env_protection == 0.0 get the
+        # guard set on the descriptor block, not just a warning
+        descriptor = DescrptSeA(
+            rcut=4.0, rcut_smth=0.5, sel=[8, 8, 0], env_protection=0.0
+        )
+        fitting = DensityFittingNet(descriptor.get_ntypes(), descriptor.get_dim_out())
+        model = DPDensityAtomicModel(descriptor, fitting, type_map=["O", "H", "X"])
+        self.assertEqual(model.descriptor.get_env_protection(), 1e-6)
+
+    def test_env_protection_hybrid_defaults(self) -> None:
+        cases = (
+            ("all omitted", [None], [1e-6]),
+            ("shared explicit value", [0.1, None], [0.1, 0.1]),
+        )
+        for case, values, expected in cases:
+            config = deepcopy(self.config)
+            base = config["model"].pop("descriptor")
+            descriptors = []
+            for value in values:
+                descriptor = deepcopy(base)
+                descriptor.pop("env_protection", None)
+                if value is not None:
+                    descriptor["env_protection"] = value
+                descriptors.append(descriptor)
+            config["model"]["descriptor"] = {
+                "type": "hybrid",
+                "list": descriptors,
+            }
+            with self.subTest(case=case):
+                normalized = normalize(config)
+                actual = [
+                    sub["env_protection"]
+                    for sub in normalized["model"]["descriptor"]["list"]
+                ]
+                self.assertEqual(actual, expected)
+
+    def test_injected_pass_honors_pair_exclude_types(self) -> None:
+        # the injected stat pass must see the model's pair_exclude_types,
+        # which are only written by the wrapped sampler
+        atomic_model = self.torch_model.atomic_model
+        atomic_model.reinit_pair_exclude([(0, 1)])
+        seen: list = []
+        original = atomic_model.descriptor.compute_input_stats
+
+        def spy(merged, path=None):
+            samples = merged() if callable(merged) else merged
+            seen.extend(samples)
+            original(merged, path)
+
+        atomic_model.descriptor.compute_input_stats = spy  # type: ignore[method-assign]
+        try:
+
+            def sampler() -> list:
+                return [self._stat_sample(atomic_model)]
+
+            atomic_model.compute_or_load_stat(sampler, stat_file_path=None)
+        finally:
+            atomic_model.descriptor.compute_input_stats = original  # type: ignore[method-assign]
+            atomic_model.reinit_pair_exclude([])
+        # injected samples carry the grid pseudo-atoms (more rows than atoms)
+        natoms = len(np.loadtxt(self.system / "type.raw", dtype=int))
+        injected = [s for s in seen if s["atype"].shape[1] > natoms and "grid" in s]
+        self.assertTrue(injected, "no injected samples reached the descriptor")
+        for sample in injected:
+            excluded = [tuple(pair) for pair in sample["pair_exclude_types"]]
+            self.assertIn((0, 1), excluded, "model exclusions not honoured")
+            self.assertIn((2, 2), excluded, "grid-grid pairs not excluded")
+
+    def test_stat_file_grid_row_writeback(self) -> None:
+        # the patched grid-type row is written back to the stat cache, and a
+        # complete cache then takes the fast path without the injected pass
+        stat_dir = DPPath(tempfile.mkdtemp(dir=self.tmpdir.name), "w")
+        atomic_model = self.torch_model.atomic_model
+
+        def sampler() -> list:
+            return [self._stat_sample(atomic_model)]
+
+        atomic_model.compute_or_load_stat(sampler, stat_file_path=stat_dir)
+        grid_type = len(self.config["model"]["type_map"]) - 1
+        r_x = list(stat_dir.rglob(f"r_{grid_type}"))
+        self.assertTrue(r_x, "no grid-type stat item written to the cache")
+        for path in r_x:
+            self.assertNotEqual(
+                float(path.load_numpy()[0]),
+                0.0,
+                f"{path} still holds zero samples (placeholder row)",
+            )
+        # second call with the complete cache: the injected pass must not run
+        original_inject = atomic_model._inject_grid_samples
+
+        def boom(_sampled: list) -> list:
+            raise AssertionError("injected pass should not run on a complete cache")
+
+        atomic_model._inject_grid_samples = boom  # type: ignore[method-assign]
+        try:
+            atomic_model.compute_or_load_stat(sampler, stat_file_path=stat_dir)
+        finally:
+            atomic_model._inject_grid_samples = original_inject  # type: ignore[method-assign]
+
+    def test_stat_cache_rejects_zero_sample_angular_grid_row(self) -> None:
+        stat_dir = DPPath(tempfile.mkdtemp(dir=self.tmpdir.name), "w")
+        atomic_model = self.torch_model.atomic_model
+
+        def sampler() -> list:
+            return [self._stat_sample(atomic_model)]
+
+        atomic_model.compute_or_load_stat(sampler, stat_file_path=stat_dir)
+        grid_type = len(self.config["model"]["type_map"]) - 1
+        angular_rows = list(stat_dir.rglob(f"a_{grid_type}"))
+        self.assertTrue(angular_rows, "no angular grid-type statistic in cache")
+        for path in angular_rows:
+            value = path.load_numpy()
+            value[0] = 0.0
+            path.save_numpy(value)
+
+        self.assertTrue(atomic_model._grid_stat_missing(stat_dir))
+
+    def test_stat_without_grid_warns_and_uses_defaults(self) -> None:
+        # Statistics can be requested before a sampler is augmented with the
+        # density input; this must take the documented fallback rather than
+        # assuming every sample contains grid coordinates.
+        atomic_model = get_model(deepcopy(model_density)).atomic_model
+        sample = self._stat_sample(atomic_model)
+        sample.pop("grid")
+        with self.assertLogs(
+            "deepmd.pt.model.atomic_model.density_atomic_model", level="WARNING"
+        ) as logs:
+            atomic_model.compute_or_load_stat(lambda: [sample])
+        self.assertIn("No grid data", "\n".join(logs.output))
+
+    def test_stat_writeback_skips_missing_cache_directory(self) -> None:
+        atomic_model = get_model(deepcopy(model_density)).atomic_model
+        grid_stats = [
+            (block["davg"][-1].detach().clone(), block["dstd"][-1].detach().clone())
+            for block in atomic_model._descriptor_stat_blocks()
+        ]
+        stat_dir = DPPath(tempfile.mkdtemp(dir=self.tmpdir.name), "w")
+        # A cache can lack one descriptor block (for example, after an older
+        # partial run); that block is deliberately skipped during write-back.
+        atomic_model._write_back_grid_stat(stat_dir, grid_stats)
+
+    def test_read_only_stat_cache_patches_grid_row_in_memory(self) -> None:
+        # Older complete caches can contain a zero-sample placeholder for the
+        # reserved grid type. Read mode may patch that row in memory, but must
+        # never attempt the write-back that update mode performs.
+        stat_file = Path(self.tmpdir.name) / "density-stat.hdf5"
+        model = get_model(deepcopy(model_density))
+        atomic_model = model.atomic_model
+
+        def sampler() -> list:
+            return [self._stat_sample(atomic_model)]
+
+        with open_stat_file(StatFileSpec(str(stat_file))) as update_path:
+            self.assertIsNotNone(update_path)
+            # The parent pass writes a complete cache with the ordinary
+            # zero-sample placeholder for the reserved grid type.
+            DPAtomicModel.compute_or_load_stat(atomic_model, sampler, update_path)
+        original = stat_file.read_bytes()
+
+        with open_stat_file(StatFileSpec(str(stat_file), "read")) as read_path:
+            self.assertIsNotNone(read_path)
+            with self.assertLogs(
+                "deepmd.pt.model.atomic_model.density_atomic_model", level="WARNING"
+            ) as logs:
+                atomic_model.compute_or_load_stat(sampler, read_path)
+        self.assertIn("read-only", "\n".join(logs.output))
+        self.assertEqual(stat_file.read_bytes(), original)
+
+    def test_grid_type_statistics(self) -> None:
+        # the reserved grid type X gets real input statistics from the
+        # injected grid samples, not the descriptor's placeholder defaults
+        descriptor = self.torch_model.atomic_model.descriptor
+        dstd = descriptor.sea["dstd"].detach().cpu().numpy()
+        self.assertEqual(dstd.shape[0], 3)
+        self.assertTrue(np.isfinite(dstd).all())
+        # placeholder default is 0.1; real statistics differ from it
+        self.assertFalse(
+            np.allclose(dstd[-1], 0.1, atol=1e-3),
+            f"grid type still has placeholder statistics: {dstd[-1]}",
+        )
+
+    def test_grid_type_statistics_dpa2(self) -> None:
+        # DPA-2 carries several stat blocks (repinit, repformers,
+        # repinit_three_body); the grid-type row must be patched in all of
+        # them, not just the first one
+        config = deepcopy(self.config)
+        config["model"]["descriptor"] = {
+            "type": "dpa2",
+            "repinit": {
+                "tebd_dim": 4,
+                "rcut": 4.0,
+                "rcut_smth": 0.5,
+                "nsel": 16,
+                "neuron": [8, 16],
+                "axis_neuron": 4,
+                "activation_function": "tanh",
+                "use_three_body": True,
+                "three_body_sel": 8,
+                "three_body_rcut": 2.0,
+                "three_body_rcut_smth": 1.0,
+            },
+            "repformer": {
+                "rcut": 2.0,
+                "rcut_smth": 1.5,
+                "nsel": 8,
+                "nlayers": 2,
+                "g1_dim": 16,
+                "g2_dim": 8,
+                "attn2_hidden": 8,
+                "attn2_nhead": 2,
+                "attn1_hidden": 16,
+                "attn1_nhead": 2,
+                "axis_neuron": 4,
+            },
+            "env_protection": 1e-6,
+            "seed": 1,
+        }
+        trainer = get_trainer(normalize(config))
+        descriptor = trainer.model.atomic_model.descriptor
+        for name in ("repinit", "repformers", "repinit_three_body"):
+            block = getattr(descriptor, name)
+            dstd = block["dstd"].detach().cpu().numpy()
+            self.assertTrue(np.isfinite(dstd).all(), name)
+            self.assertFalse(
+                np.allclose(dstd[-1], 0.1, atol=1e-3),
+                f"{name} still has placeholder statistics for the grid type",
+            )
+        # the reserved grid type X gets real input statistics from the
+        # injected grid samples, not the descriptor's placeholder defaults
+        descriptor = self.torch_model.atomic_model.descriptor
+        dstd = descriptor.sea["dstd"].detach().cpu().numpy()
+        self.assertEqual(dstd.shape[0], 3)
+        self.assertTrue(np.isfinite(dstd).all())
+        # placeholder default is 0.1; real statistics differ from it
+        self.assertFalse(
+            np.allclose(dstd[-1], 0.1, atol=1e-3),
+            f"grid type still has placeholder statistics: {dstd[-1]}",
+        )
+
+    def test_eval_requires_grid(self) -> None:
+        # evaluating a density model without grid must fail with a clear
+        # error instead of a bare KeyError on the output name table
+        dp = DeepDensity(self.model_path)
+        set_dir = self.system / "set.000"
+        coord = np.load(set_dir / "coord.npy")[:2]
+        box = np.load(set_dir / "box.npy")[:2]
+        atype = np.loadtxt(self.system / "type.raw", dtype=int)
+        with self.assertRaisesRegex(ValueError, "grid is required"):
+            dp.deep_eval.eval(coord, box, atype)
+
+    def test_eval_no_auto_batch(self) -> None:
+        # auto_batch_size=False bypasses execute_all's single-tuple unwrapping;
+        # the density result must still be normalized to a bare ndarray
+        dp = DeepDensity(self.model_path, auto_batch_size=False)
+        set_dir = self.system / "set.000"
+        coord = np.load(set_dir / "coord.npy")[:2]
+        box = np.load(set_dir / "box.npy")[:2]
+        grid = np.load(set_dir / "grid.npy")[:2]
+        atype = np.loadtxt(self.system / "type.raw", dtype=int)
+        out = dp.eval(coord, box, atype, grid=grid)
+        self.assertEqual(out.shape, (2, self.ngrid))
+        # and match the default-batching result numerically
+        out_default = DeepDensity(self.model_path).eval(coord, box, atype, grid=grid)
+        np.testing.assert_allclose(out, out_default)
+
+    def test_eval_without_periodic_box(self) -> None:
+        set_dir = self.system / "set.000"
+        coord = np.load(set_dir / "coord.npy")[:1]
+        grid = np.load(set_dir / "grid.npy")[:1]
+        atype = np.loadtxt(self.system / "type.raw", dtype=int)
+        out = DeepDensity(self.model_path).eval(coord, None, atype, grid=grid)
+        self.assertEqual(out.shape, (1, self.ngrid))
+        self.assertTrue(np.isfinite(out).all())
+
+    def test_grid_at_atoms_finite(self) -> None:
+        # grid points coincident with atoms are legitimate inputs; with
+        # env_protection > 0 (set in the test config) the predictions and
+        # gradients must stay finite instead of turning NaN via 1/r terms
+        set_dir = self.system / "set.000"
+        coord = np.load(set_dir / "coord.npy")[:2]
+        box = np.load(set_dir / "box.npy")[:2]
+        atype = np.loadtxt(self.system / "type.raw", dtype=int)
+        grid_at_atoms = coord.reshape(coord.shape[0], -1, 3)
+        out = DeepDensity(self.model_path).eval(coord, box, atype, grid=grid_at_atoms)
+        self.assertTrue(np.isfinite(out).all())
+
+        # gradients wrt the grid input and all parameters stay finite
+        input_dict = {
+            kk: vv.clone() if isinstance(vv, torch.Tensor) else vv
+            for kk, vv in self.input_dict.items()
+        }
+        coord_t = input_dict["coord"]
+        grid_t = (
+            coord_t.reshape(coord_t.shape[0], -1, 3).detach().clone().requires_grad_()
+        )
+        input_dict["grid"] = grid_t
+        model_out = self.torch_model(**input_dict)
+        self.assertTrue(torch.isfinite(model_out["density"]).all())
+        model_out["density"].sum().backward()
+        self.assertIsNotNone(grid_t.grad)
+        self.assertTrue(torch.isfinite(grid_t.grad).all())
+        for param in self.torch_model.parameters():
+            if param.grad is not None:
+                self.assertTrue(torch.isfinite(param.grad).all())
+
+    def test_grid_periodic_translation(self) -> None:
+        # grids shifted by integer cell vectors are periodically equivalent
+        # and must produce the same density as the wrapped originals
+        set_dir = self.system / "set.000"
+        coord = np.load(set_dir / "coord.npy")[:2]
+        box = np.load(set_dir / "box.npy")[:2]
+        grid = np.load(set_dir / "grid.npy")[:2]
+        atype = np.loadtxt(self.system / "type.raw", dtype=int)
+        dp = DeepDensity(self.model_path)
+        out = dp.eval(coord, box, atype, grid=grid)
+        cell_vec = box.reshape(box.shape[0], 3, 3)[:, 0]
+        out_shift = dp.eval(coord, box, atype, grid=grid + 2.0 * cell_vec[:, None, :])
+        np.testing.assert_allclose(out_shift, out, rtol=1e-5, atol=1e-5)
+
+    def test_dp_test_without_detail_file(self) -> None:
+        dp = DeepDensity(self.model_path)
+        data = DeepmdData(
+            str(self.system),
+            set_prefix="set",
+            shuffle_test=False,
+            type_map=dp.get_type_map(),
+            sort_atoms=False,
+        )
+        errors = build_tester(dp, atomic=False).run(
+            data,
+            str(self.system),
+            numb_test=self.nframes,
+            detail_file=None,
+        )
+        self.assertEqual(set(errors), {"mae_density", "rmse_density"})
+
+    def test_dp_test_chunking_matches_whole_run(self) -> None:
+        def run(chunk_size: int, detail_file: str) -> dict:
+            dp = DeepDensity(self.model_path)
+            data = DeepmdData(
+                str(self.system),
+                set_prefix="set",
+                shuffle_test=False,
+                type_map=dp.get_type_map(),
+                sort_atoms=False,
+            )
+            with patch.dict(
+                os.environ, {"DP_TEST_CHUNK_ATOMS": str(chunk_size)}, clear=False
+            ):
+                return build_tester(dp, atomic=False).run(
+                    data,
+                    str(self.system),
+                    numb_test=self.nframes,
+                    detail_file=detail_file,
+                )
+
+        whole_detail = os.path.join(self.tmpdir.name, "detail_whole")
+        chunk_detail = os.path.join(self.tmpdir.name, "detail_chunk")
+        whole = run(1_000_000, whole_detail)
+        chunked = run(self.ngrid, chunk_detail)
+        self.assertEqual(whole.keys(), chunked.keys())
+        for key in whole:
+            np.testing.assert_allclose(whole[key], chunked[key])
+        for frame in range(self.nframes):
+            np.testing.assert_allclose(
+                np.loadtxt(f"{whole_detail}.density.out.{frame}"),
+                np.loadtxt(f"{chunk_detail}.density.out.{frame}"),
+            )
+
+    def test_dp_test_mixed_type(self) -> None:
+        system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "mixed-system"
+        shutil.copytree(self.system, system)
+        atom_types = np.array(
+            [
+                [0, 0, 1, 1, 1],
+                [0, 1, 0, 1, 1],
+                [1, 0, 0, 1, 1],
+                [1, 1, 0, 0, 1],
+            ],
+            dtype=np.int32,
+        )
+        np.save(system / "set.000" / "real_atom_types.npy", atom_types)
+        detail_file = os.path.join(self.tmpdir.name, "detail_mixed")
+        dp_test(
+            model=self.model_path,
+            system=str(system),
+            datafile=None,
+            numb_test=0,
+            rand_seed=None,
+            shuffle_test=False,
+            detail_file=detail_file,
+            atomic=False,
+        )
+        set_dir = system / "set.000"
+        ref = DeepDensity(self.model_path).eval(
+            np.load(set_dir / "coord.npy"),
+            np.load(set_dir / "box.npy"),
+            atom_types,
+            grid=np.load(set_dir / "grid.npy"),
+            mixed_type=True,
+        )
+        for frame in range(self.nframes):
+            detail = np.loadtxt(f"{detail_file}.density.out.{frame}", skiprows=1)
+            np.testing.assert_allclose(detail[:, 1], ref[frame], rtol=1e-5, atol=1e-6)
+
+    def test_dp_test_forwards_fparam(self) -> None:
+        params = deepcopy(model_density)
+        params["fitting_net"]["numb_fparam"] = 1
+        model = get_model(params)
+        fd, model_path = tempfile.mkstemp(suffix=".pth")
+        os.close(fd)
+        system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "fparam-system"
+        shutil.copytree(self.system, system)
+        fparam = np.arange(1, self.nframes + 1, dtype=np.float64).reshape(-1, 1)
+        np.save(system / "set.000" / "fparam.npy", fparam)
+        try:
+            torch.jit.save(torch.jit.script(model), model_path)
+            detail_file = os.path.join(self.tmpdir.name, "detail_fparam")
+            dp_test(
+                model=model_path,
+                system=str(system),
+                datafile=None,
+                numb_test=0,
+                rand_seed=None,
+                shuffle_test=False,
+                detail_file=detail_file,
+                atomic=False,
+            )
+            set_dir = system / "set.000"
+            dp = DeepDensity(model_path)
+            args = (
+                np.load(set_dir / "coord.npy"),
+                np.load(set_dir / "box.npy"),
+                np.loadtxt(system / "type.raw", dtype=int),
+            )
+            grid = np.load(set_dir / "grid.npy")
+            ref = dp.eval(*args, grid=grid, fparam=fparam)
+            without_fparam_effect = dp.eval(
+                *args, grid=grid, fparam=np.zeros_like(fparam)
+            )
+            self.assertFalse(np.allclose(ref, without_fparam_effect))
+            for frame in range(self.nframes):
+                detail = np.loadtxt(f"{detail_file}.density.out.{frame}", skiprows=1)
+                np.testing.assert_allclose(
+                    detail[:, 1], ref[frame], rtol=1e-5, atol=1e-6
+                )
+        finally:
+            os.unlink(model_path)
+
+    def test_dp_test_uses_default_fparam(self) -> None:
+        params = deepcopy(model_density)
+        params["fitting_net"].update(numb_fparam=1, default_fparam=[0.5])
+        fd, model_path = tempfile.mkstemp(suffix=".pth")
+        os.close(fd)
+        default_system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "default"
+        explicit_system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "explicit"
+        shutil.copytree(self.system, default_system)
+        shutil.copytree(self.system, explicit_system)
+        np.save(
+            explicit_system / "set.000" / "fparam.npy",
+            np.full((self.nframes, 1), 0.5),
+        )
+        try:
+            torch.jit.save(torch.jit.script(get_model(params)), model_path)
+            detail_paths = []
+            for name, system in (
+                ("default", default_system),
+                ("explicit", explicit_system),
+            ):
+                detail_file = os.path.join(self.tmpdir.name, f"detail_fparam_{name}")
+                dp_test(
+                    model=model_path,
+                    system=str(system),
+                    datafile=None,
+                    numb_test=0,
+                    rand_seed=None,
+                    shuffle_test=False,
+                    detail_file=detail_file,
+                    atomic=False,
+                )
+                detail_paths.append(detail_file)
+            for frame in range(self.nframes):
+                np.testing.assert_allclose(
+                    np.loadtxt(f"{detail_paths[0]}.density.out.{frame}"),
+                    np.loadtxt(f"{detail_paths[1]}.density.out.{frame}"),
+                )
+        finally:
+            os.unlink(model_path)
+
+    def test_dp_test_without_periodic_box(self) -> None:
+        system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "nopbc-system"
+        shutil.copytree(self.system, system)
+        (system / "set.000" / "box.npy").unlink()
+        (system / "nopbc").touch()
+        detail_file = os.path.join(self.tmpdir.name, "detail_nopbc")
+        dp_test(
+            model=self.model_path,
+            system=str(system),
+            datafile=None,
+            numb_test=0,
+            rand_seed=None,
+            shuffle_test=False,
+            detail_file=detail_file,
+            atomic=False,
+        )
+        set_dir = system / "set.000"
+        ref = DeepDensity(self.model_path).eval(
+            np.load(set_dir / "coord.npy"),
+            None,
+            np.loadtxt(system / "type.raw", dtype=int),
+            grid=np.load(set_dir / "grid.npy"),
+        )
+        for frame in range(self.nframes):
+            detail = np.loadtxt(f"{detail_file}.density.out.{frame}", skiprows=1)
+            np.testing.assert_allclose(detail[:, 1], ref[frame], rtol=1e-5, atol=1e-6)
+
+    def test_dp_test_shuffle(self) -> None:
+        # grid/density must be shuffled together with the frames: every
+        # (label, pred) row written by the shuffled run must match a row
+        # from the unshuffled reference run
+        detail_ref = os.path.join(self.tmpdir.name, "detail_ref")
+        dp_test(
+            model=self.model_path,
+            system=str(self.system),
+            datafile=None,
+            numb_test=0,
+            rand_seed=None,
+            shuffle_test=False,
+            detail_file=detail_ref,
+            atomic=False,
+        )
+        ref_rows = set()
+        for frame in range(self.nframes):
+            detail = np.loadtxt(f"{detail_ref}.density.out.{frame}", skiprows=1)
+            for row in detail:
+                ref_rows.add(tuple(np.round(row, decimals=5)))
+
+        detail_shuf = os.path.join(self.tmpdir.name, "detail_shuf")
+        dp_test(
+            model=self.model_path,
+            system=str(self.system),
+            datafile=None,
+            numb_test=2,
+            rand_seed=42,
+            shuffle_test=True,
+            detail_file=detail_shuf,
+            atomic=False,
+        )
+        shuf_rows = []
+        for frame in range(2):
+            detail = np.loadtxt(f"{detail_shuf}.density.out.{frame}", skiprows=1)
+            shuf_rows.extend(tuple(np.round(row, decimals=5)) for row in detail)
+        self.assertEqual(len(shuf_rows), 2 * self.ngrid)
+        for row in shuf_rows:
+            self.assertIn(row, ref_rows)
+
+    def test_dp_test_flattened_grid(self) -> None:
+        # the loader also accepts a flattened (nframes, ngrid*3) grid; the
+        # tester must normalize it to the (nframes, ngrid, 3) contract that
+        # DeepDensity.eval expects, so both layouts give the same details
+        system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "system"
+        shutil.copytree(self.system, system)
+        set_dir = system / "set.000"
+        grid3d = np.load(set_dir / "grid.npy")
+        np.save(set_dir / "grid.npy", grid3d.reshape(self.nframes, -1))
+        detail_file = os.path.join(self.tmpdir.name, "detail_flat")
+        dp_test(
+            model=self.model_path,
+            system=str(system),
+            datafile=None,
+            numb_test=0,
+            rand_seed=None,
+            shuffle_test=False,
+            detail_file=detail_file,
+            atomic=False,
+        )
+        coord = np.load(set_dir / "coord.npy")
+        box = np.load(set_dir / "box.npy")
+        atype = np.loadtxt(system / "type.raw", dtype=int)
+        ref = DeepDensity(self.model_path).eval(coord, box, atype, grid=grid3d)
+        for frame in range(self.nframes):
+            detail = np.loadtxt(f"{detail_file}.density.out.{frame}", skiprows=1)
+            self.assertEqual(detail.shape, (self.ngrid, 2))
+            np.testing.assert_allclose(detail[:, 1], ref[frame], rtol=1e-5, atol=1e-6)
+
+
+model_density_dpa3_chg_spin = {
+    "type_map": ["O", "H", "X"],
+    "descriptor": {
+        "type": "dpa3",
+        "repflow": {
+            "n_dim": 8,
+            "e_dim": 6,
+            "a_dim": 4,
+            "nlayers": 1,
+            "e_rcut": 4.0,
+            "e_rcut_smth": 0.5,
+            "e_sel": 12,
+            "a_rcut": 3.5,
+            "a_rcut_smth": 0.5,
+            "a_sel": 8,
+            "axis_neuron": 4,
+            "update_angle": False,
+        },
+        "env_protection": 1e-6,
+        "add_chg_spin_ebd": True,
+        "default_chg_spin": [1.0, 0.0],
+        "seed": 1,
+    },
+    "fitting_net": {
+        "type": "density",
+        "neuron": [8, 8],
+        "seed": 1,
+    },
+}
+
+
+class TestDPTestDensityChgSpin(unittest.TestCase):
+    """charge_spin forwarding on the density inference and dp test paths."""
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.tmpdir = tempfile.TemporaryDirectory()
+        root = Path(cls.tmpdir.name)
+        rng = np.random.default_rng(7)
+        cls.nframes, natoms, cls.ngrid = 4, 5, 8
+        cls.system = root / "system"
+        set_dir = cls.system / "set.000"
+        set_dir.mkdir(parents=True)
+        np.save(set_dir / "coord.npy", rng.random((cls.nframes, natoms * 3)) * 8.0)
+        np.save(
+            set_dir / "box.npy",
+            np.tile(np.eye(3).reshape(-1) * 10.0, (cls.nframes, 1)),
+        )
+        np.save(set_dir / "grid.npy", rng.random((cls.nframes, cls.ngrid, 3)) * 8.0)
+        np.save(set_dir / "density.npy", rng.random((cls.nframes, cls.ngrid, 1)))
+        # non-default charge/spin conditions, so forwarding has a visible effect
+        np.save(set_dir / "charge_spin.npy", np.tile([2.0, 1.0], (cls.nframes, 1)))
+        np.savetxt(cls.system / "type.raw", [0, 0, 1, 1, 1], fmt="%d")
+        np.savetxt(cls.system / "type_map.raw", ["O", "H", "X"], fmt="%s")
+
+        model = get_model(deepcopy(model_density_dpa3_chg_spin))
+        scripted = torch.jit.script(model)
+        tmp_fd, cls.model_path = tempfile.mkstemp(suffix=".pth")
+        os.close(tmp_fd)
+        torch.jit.save(scripted, cls.model_path)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        os.unlink(cls.model_path)
+        cls.tmpdir.cleanup()
+
+    def _inputs(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        set_dir = self.system / "set.000"
+        coord = np.load(set_dir / "coord.npy")
+        box = np.load(set_dir / "box.npy")
+        grid = np.load(set_dir / "grid.npy")
+        atype = np.loadtxt(self.system / "type.raw", dtype=int)
+        return coord, box, grid, atype
+
+    def test_eval_forwards_charge_spin(self) -> None:
+        dp = DeepDensity(self.model_path)
+        self.assertTrue(dp.has_chg_spin_ebd())
+        self.assertTrue(dp.has_default_chg_spin())
+        coord, box, grid, atype = self._inputs()
+        # without explicit charge_spin the model falls back to its default
+        out_default = dp.eval(coord, box, atype, grid=grid)
+        out_explicit_default = dp.eval(
+            coord,
+            box,
+            atype,
+            grid=grid,
+            charge_spin=np.tile([1.0, 0.0], (self.nframes, 1)),
+        )
+        np.testing.assert_allclose(
+            out_default, out_explicit_default, rtol=1e-6, atol=1e-8
+        )
+        # a non-default condition must change the prediction: the value has
+        # to reach the descriptor, not be dropped on the eval path
+        out_other = dp.eval(
+            coord,
+            box,
+            atype,
+            grid=grid,
+            charge_spin=np.tile([2.0, 1.0], (self.nframes, 1)),
+        )
+        self.assertFalse(np.allclose(out_default, out_other))
+
+    def test_eval_without_charge_spin_or_default(self) -> None:
+        # A charge/spin-conditioned descriptor needs either an explicit
+        # condition or a configured default; fail clearly before the
+        # descriptor's internal assertion.
+        config = deepcopy(model_density_dpa3_chg_spin)
+        config["descriptor"].pop("default_chg_spin")
+        path = os.path.join(self.tmpdir.name, "no-default-charge-spin.pth")
+        with torch.device("cpu"):
+            torch.jit.save(torch.jit.script(get_model(config)), path)
+        coord, box, grid, atype = self._inputs()
+        with self.assertRaisesRegex(torch.jit.Error, "charge_spin is required"):
+            DeepDensity(path).eval(coord, box, atype, grid=grid)
+
+    def test_eval_grid_chunking_dpa3_matches_full_grid(self) -> None:
+        coord, box, grid, atype = self._inputs()
+        dp = DeepDensity(self.model_path)
+
+        full = dp.eval(coord, box, atype, grid=grid, grid_batch_size=None)
+        chunked = dp.eval(coord, box, atype, grid=grid, grid_batch_size=3)
+
+        np.testing.assert_allclose(chunked, full, rtol=1e-6, atol=1e-8)
+
+    def test_dp_test_charge_spin_data_and_default(self) -> None:
+        coord, box, grid, atype = self._inputs()
+        dp = DeepDensity(self.model_path)
+        explicit_ref = dp.eval(
+            coord,
+            box,
+            atype,
+            grid=grid,
+            charge_spin=np.tile([2.0, 1.0], (self.nframes, 1)),
+        )
+        default_ref = dp.eval(coord, box, atype, grid=grid)
+        self.assertFalse(
+            np.allclose(explicit_ref, default_ref),
+            "the test is vacuous: charge_spin has no effect on this model",
+        )
+
+        system_without_data = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "system"
+        shutil.copytree(self.system, system_without_data)
+        (system_without_data / "set.000" / "charge_spin.npy").unlink()
+        cases = (
+            ("explicit data", self.system, explicit_ref),
+            ("model default", system_without_data, default_ref),
+        )
+        for case, system, reference in cases:
+            detail_file = os.path.join(self.tmpdir.name, f"detail_cs_{case}")
+            with self.subTest(case=case):
+                dp_test(
+                    model=self.model_path,
+                    system=str(system),
+                    datafile=None,
+                    numb_test=0,
+                    rand_seed=None,
+                    shuffle_test=False,
+                    detail_file=detail_file,
+                    atomic=False,
+                )
+                for frame in range(self.nframes):
+                    detail = np.loadtxt(
+                        f"{detail_file}.density.out.{frame}", skiprows=1
+                    )
+                    self.assertEqual(detail.shape, (self.ngrid, 2))
+                    np.testing.assert_allclose(
+                        detail[:, 1], reference[frame], rtol=1e-5, atol=1e-6
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/test_dp_test_spin_virtual_atoms.py b/source/tests/pt/test_dp_test_spin_virtual_atoms.py
index 5318cb39e6..dba6d4c092 100644
--- a/source/tests/pt/test_dp_test_spin_virtual_atoms.py
+++ b/source/tests/pt/test_dp_test_spin_virtual_atoms.py
@@ -147,9 +147,10 @@ def test_dp_test_spin_virtual_atom_metrics(tmp_path, monkeypatch, padding_label)
         capture_output=True,
         text=True,
         timeout=120,
-        check=True,
+        check=False,
     )
     output = result.stdout + result.stderr
+    assert result.returncode == 0, output
     for label, key in (("Force atom RMSE", "rmse_fr"), ("Force spin RMSE", "rmse_fm")):
         assert label in output
         assert f"{expected[key][0]:e}" in output
diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py
index ed27015a92..fc7352f0b8 100644
--- a/source/tests/pt/test_lmdb_dataloader.py
+++ b/source/tests/pt/test_lmdb_dataloader.py
@@ -6,12 +6,28 @@
 Consistency tests (dpmodel vs pt) live in source/tests/consistent/test_lmdb_data.py.
 """
 
+import gc
+from itertools import (
+    chain,
+    islice,
+    repeat,
+)
+from pathlib import (
+    Path,
+)
+from typing import (
+    Any,
+)
+
 import lmdb
 import msgpack
 import numpy as np
 import pytest
 import torch
 
+from deepmd.dpmodel.train import (
+    resolve_step_schedule,
+)
 from deepmd.dpmodel.utils import (
     lmdb_data,
 )
@@ -30,6 +46,10 @@
 from deepmd.pt.loss.ener import (
     EnergyStdLoss,
 )
+from deepmd.pt.utils.dataloader import (
+    create_dataloader,
+    create_dataset,
+)
 from deepmd.pt.utils.lmdb_dataset import (
     LmdbBatchDataLoader,
     LmdbDataset,
@@ -38,9 +58,22 @@
 from deepmd.pt.utils.stat import (
     make_stat_input,
 )
+from deepmd.pt_expt.utils.lmdb_dataset import (
+    LmdbDataSystem,
+)
 from deepmd.utils.data import (
     DataRequirementItem,
 )
+from deepmd.utils.data_group import (
+    iter_data_sources,
+)
+from deepmd.utils.data_system import (
+    get_data,
+)
+
+from ..common.dpmodel.test_lmdb_data import (
+    _create_mixed_sid_nloc_lmdb,
+)
 
 
 def _make_frame(natoms: int = 6, seed: int = 0) -> dict:
@@ -1043,6 +1076,31 @@ def auto_prob_lmdb(tmp_path):
 class TestAutoProbDataset:
     """Test LmdbDataset with auto_prob_style."""
 
+    @pytest.mark.parametrize(
+        ("auto_prob", "expected_frames"),
+        [("prob_uniform", 300), ("prob_sys_size;0:1:0.5;1:2:0.5", 350)],
+    )
+    def test_summary_reports_exact_pass_frames(
+        self,
+        auto_prob_lmdb: str,
+        caplog: pytest.LogCaptureFixture,
+        auto_prob: str,
+        expected_frames: int,
+    ) -> None:
+        ds = LmdbDataset(
+            auto_prob_lmdb,
+            type_map=["O", "H"],
+            batch_size=4,
+            auto_prob_style=auto_prob,
+        )
+        try:
+            with caplog.at_level("INFO", logger="deepmd.pt.utils.lmdb_dataset"):
+                ds.print_summary("training", None)
+            assert f"300->{expected_frames} frames" in caplog.text
+            assert sum(len(batch) for batch in ds._batch_sampler) == expected_frames
+        finally:
+            ds.close()
+
     def test_dataset_auto_prob_passthrough(self, auto_prob_lmdb):
         ds = LmdbDataset(
             auto_prob_lmdb,
@@ -1284,6 +1342,47 @@ def test_merge_rejects_mixed_explicit_and_missing_type_maps(self, tmp_path):
         assert "missing" in str(exc.value)
         assert not dst.exists()
 
+    def test_merge_keeps_cached_source_readers_open(self, tmp_path):
+        """A merge must release its cache reference without closing readers."""
+        src = str(tmp_path / "shared_src.lmdb")
+        dst = str(tmp_path / "shared_dst.lmdb")
+        _create_lmdb_with_system_ids(
+            src, system_frames=[3], natoms=6, type_map=["O", "H"]
+        )
+        cache_key = str(tmp_path.joinpath("shared_src.lmdb").resolve())
+        existing_reader = LmdbDataReader(src, ["O", "H"])
+        new_reader = None
+        expected_coord = existing_reader[0]["coord"].copy()
+        cached_env, refcount_before_merge = lmdb_data._ENV_CACHE[cache_key]
+
+        try:
+            merge_lmdb([src], dst)
+
+            # merge_lmdb shares the process-level environment cache with readers.
+            # Its temporary reference must be balanced without replacing or
+            # invalidating the environment retained by the existing reader.
+            env_after_merge, refcount_after_merge = lmdb_data._ENV_CACHE[cache_key]
+            assert env_after_merge is cached_env
+            assert refcount_after_merge == refcount_before_merge
+            np.testing.assert_array_equal(existing_reader[0]["coord"], expected_coord)
+
+            new_reader = LmdbDataReader(src, ["O", "H"])
+            env_after_new_reader, refcount_after_new_reader = lmdb_data._ENV_CACHE[
+                cache_key
+            ]
+            assert env_after_new_reader is cached_env
+            assert refcount_after_new_reader == refcount_before_merge + 1
+            np.testing.assert_array_equal(new_reader[0]["coord"], expected_coord)
+        finally:
+            # Drop both strong references so their finalizers release the
+            # cache entries before the cleanup assertion below.
+            del existing_reader, new_reader
+            gc.collect()
+            # Keep a failing regression isolated from later LMDB tests. On the
+            # fixed path reader finalizers normally remove this entry already.
+            while cache_key in lmdb_data._ENV_CACHE:
+                lmdb_data._close_lmdb(src)
+
 
 # ============================================================
 # Multitask LMDB training
@@ -1491,7 +1590,13 @@ class TestMultitaskLmdbTraining:
     All assertions are in a single test to avoid creating multiple trainers.
     """
 
-    def test_multitask_lmdb_end_to_end(self, multitask_lmdb_setup, monkeypatch):
+    @pytest.mark.parametrize("grouped", [False, True])
+    def test_multitask_lmdb_end_to_end(
+        self,
+        multitask_lmdb_setup: tuple[dict[str, Any], Path],
+        monkeypatch: pytest.MonkeyPatch,
+        grouped: bool,
+    ) -> None:
         from copy import (
             deepcopy,
         )
@@ -1510,6 +1615,20 @@ def test_multitask_lmdb_end_to_end(self, multitask_lmdb_setup, monkeypatch):
         )
 
         config, tmp_path = multitask_lmdb_setup
+        if grouped:
+            branches = config["training"]["data_dict"]
+            branches["model_1"]["training_data"] = [
+                {
+                    **branches["model_1"]["training_data"],
+                    "prob": 0.6,
+                    "min_pair_dist": 0.1,
+                },
+                {**branches["model_2"]["training_data"], "prob": 0.4},
+            ]
+            branches["model_1"]["validation_data"]["systems"] = [
+                branches["model_1"]["validation_data"]["systems"],
+                branches["model_2"]["validation_data"]["systems"],
+            ]
         monkeypatch.chdir(tmp_path)
         monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "2")
         config = update_deepmd_input(deepcopy(config), warning=True)
@@ -1536,21 +1655,433 @@ def test_multitask_lmdb_end_to_end(self, multitask_lmdb_setup, monkeypatch):
             )
             assert "coord" in input_dict
             assert "sid" in log_dict
-        assert (
-            trainer.training_dataloader["model_1"]._batch_iterator._pool
-            is trainer.training_dataloader["model_2"]._batch_iterator._pool
-        )
+        native_loaders = [
+            source.loader
+            for loader in trainer.training_dataloader.values()
+            for source in iter_data_sources(loader)
+        ]
+        for loader in native_loaders:
+            next(iter(loader))
+        assert len({id(loader._batch_iterator._pool) for loader in native_loaders}) == 1
 
         # -- training run assertions --
         trainer.run()
         assert len(list(tmp_path.glob("model.ckpt*.pt"))) > 0
-        assert trainer.training_dataloader["model_1"]._batch_iterator.closed
-        assert trainer.training_dataloader["model_2"]._batch_iterator.closed
-        assert trainer.training_dataloader["model_1"].dataset._reader.closed
-        assert trainer.training_dataloader["model_2"].dataset._reader.closed
+        assert all(loader._batch_iterator.closed for loader in native_loaders)
+        assert all(loader.dataset._reader.closed for loader in native_loaders)
 
         # Explicit cleanup to free memory on CI
         import gc
 
         del trainer
         gc.collect()
+
+
+def test_multiple_lmdb_loader_uses_native_batch_sizes(
+    lmdb_dir: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    parent = tmp_path / "nested"
+    parent.mkdir()
+    second = parent / "second.lmdb"
+    _create_test_lmdb(str(second), nframes=4, natoms=8)
+    config = {"systems": [lmdb_dir, str(parent)], "batch_size": 2}
+    data = create_dataset(config, ["O", "H"])
+    loader = create_dataloader(data, config, seed=0)
+    try:
+        np.testing.assert_allclose(loader.probabilities, [5 / 7, 2 / 7])
+        assert loader.epoch_length == 7
+        batches = [batch for _ in range(3) for batch in loader]
+        assert {int(batch["sid"]) for batch in batches} == {0, 1}
+        assert {batch["coord"].shape[1] for batch in batches} == {6, 8}
+    finally:
+        loader.close()
+    assert all(source._reader.closed for source in iter_data_sources(data))
+
+
+def test_filter_can_remove_the_numpy_subset_of_a_mixed_group(
+    lmdb_dir: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    numpy = str(Path(__file__).parent / "water" / "data" / "single")
+    config = {"systems": [numpy, lmdb_dir], "batch_size": "filter:8"}
+    pt_data = create_dataset(config, ["O", "H"])
+    numpy_data = get_data(
+        config,
+        0,
+        ["O", "H"],
+        None,
+        lmdb_factory=lambda path, batch_size, auto_prob: LmdbDataSystem(
+            path, ["O", "H"], batch_size, auto_prob, num_workers=0
+        ),
+    )
+    try:
+        assert isinstance(pt_data, LmdbDataset)
+        assert numpy_data.get_batch()["coord"].shape == (1, 6, 3)
+    finally:
+        pt_data.close()
+        numpy_data.close()
+
+
+def test_group_probability_preserves_single_lmdb_auto_prob(
+    auto_prob_lmdb: str, lmdb_dir: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    config = [
+        {
+            "systems": auto_prob_lmdb,
+            "batch_size": 2,
+            "auto_prob": "prob_sys_size;0:1:0.5;1:3:0.5",
+            "prob": 0.6,
+        },
+        {"systems": lmdb_dir, "batch_size": 2, "prob": 0.4},
+    ]
+    data = create_dataset(config, ["O", "H"])
+    loader = create_dataloader(data, config, seed=0)
+    try:
+        assert data.children[0]._block_targets
+        np.testing.assert_allclose(loader.probabilities, [0.6, 0.4])
+        assert len(data.dataloaders) == 2
+    finally:
+        loader.close()
+
+
+def test_distributed_lmdb_groups_separate_mass_from_local_padding(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    paths = [str(tmp_path / f"source_{index}.lmdb") for index in range(2)]
+    for path, frames in zip(paths, (1, 3), strict=True):
+        _create_test_lmdb(path, nframes=frames)
+    config = {"systems": paths, "batch_size": 1}
+    for rank in (0, 1):
+        data = create_dataset(config, ["O", "H"])
+        loader = create_dataloader(data, config, seed=0, rank=rank, world_size=2)
+        try:
+            np.testing.assert_allclose(loader.probabilities, [0.25, 0.75])
+            np.testing.assert_array_equal(loader.batch_counts, [1, 2])
+            np.testing.assert_array_equal(loader.sampling_sizes, [1, 3])
+            assert loader.epoch_length == 4
+        finally:
+            loader.close()
+
+
+@pytest.fixture
+def uniform_sources(tmp_path: Path) -> tuple[list[str], list[np.ndarray]]:
+    paths = [str(tmp_path / f"source_{index}.lmdb") for index in range(3)]
+    _create_test_lmdb(paths[0], nframes=2, natoms=6)
+    _create_mixed_sid_nloc_lmdb(paths[1], [(1, 6), (7, 10), (3, 16)])
+    _create_test_lmdb(paths[2], nframes=2, natoms=18)
+    markers = [np.zeros(2, dtype=int), np.repeat([1, 2, 3], [1, 7, 3]), np.full(2, 4)]
+    for path, values in zip(paths, markers, strict=True):
+        with lmdb.open(path) as database, database.begin(write=True) as transaction:
+            for index, value in enumerate(values):
+                key = format(index, "012d").encode()
+                frame = msgpack.unpackb(transaction.get(key), raw=False)
+                frame["energies"] = {
+                    "type": " None:
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    paths, _ = uniform_sources
+    config = {
+        "systems": str(tmp_path) if recursive else paths,
+        "batch_size": batch_size,
+        "auto_prob": "prob_uniform",
+    }
+    requirements = [
+        DataRequirementItem("energy", 1, atomic=False, must=True, high_prec=True),
+        DataRequirementItem("virial", 9, atomic=False, must=False, high_prec=True),
+    ]
+    if backend == "pt":
+        data = create_dataset(config, ["O", "H"])
+        data.add_data_requirement(requirements)
+        owner = create_dataloader(data, config, seed=31)
+        batches = chain.from_iterable(repeat(owner))
+    else:
+        data = get_data(
+            config,
+            0,
+            ["O", "H"],
+            None,
+            seed=31,
+            lmdb_factory=lambda path, size, auto_prob: LmdbDataSystem(
+                path, ["O", "H"], size, auto_prob, seed=31, num_workers=0
+            ),
+        )
+        data.add_data_requirements(requirements)
+        data.use_ragged_batches(ragged)
+        owner = data
+        batches = (data.get_batch() for _ in repeat(None))
+    counts = np.zeros(5, dtype=np.int64)
+    mixed = False
+    try:
+        for batch in islice(batches, 2500):
+            labels = np.asarray(batch["energy"]).reshape(-1).astype(int)
+            available = np.asarray(batch["fid"]).reshape(-1).astype(int) % 2 == 0
+            assert np.all(available == available[0])
+            assert bool(batch["find_virial"]) == bool(available[0])
+            counts += np.bincount(labels, minlength=5)
+            mixed |= np.unique(labels).size > 1
+    finally:
+        owner.close()
+    eligible = 3 if batch_size == "filter:12" else 5
+    np.testing.assert_allclose(
+        counts[:eligible] / counts.sum(), 1 / eligible, atol=0.025
+    )
+    assert counts[eligible:].sum() == 0
+    if batch_size == "mix:50":
+        assert mixed
+
+
+@pytest.mark.parametrize("num_workers", [0, 1])
+def test_uniform_numpy_keeps_native_cursors_across_loader_passes(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, num_workers: int
+) -> None:
+    from deepmd.pt.utils import (
+        env,
+    )
+
+    monkeypatch.setattr(env, "NUM_WORKERS", num_workers)
+    np.savetxt(tmp_path / "type.raw", [0], fmt="%d")
+    set_path = tmp_path / "set.000"
+    set_path.mkdir()
+    np.save(set_path / "coord.npy", np.arange(24, dtype=float).reshape(8, 3))
+    np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (8, 1)))
+    config = {
+        "systems": [str(tmp_path), str(tmp_path)],
+        "batch_size": [7, 8],
+        "auto_prob": "prob_uniform",
+    }
+    data = create_dataset(config, ["O"])
+    loader = create_dataloader(data, config, seed=7)
+    try:
+        np.testing.assert_allclose(loader.probabilities, [2 / 3, 1 / 3])
+        np.testing.assert_allclose(
+            loader.probabilities * data.mean_batch_nframes, [8 / 3, 8 / 3]
+        )
+        # Deterministic source selection exposes short tails across outer passes.
+        loader.loader.sampler.weights[:] = torch.tensor(
+            [1.0, 0.0], dtype=torch.float64, device="cpu"
+        )
+        sizes = [batch["coord"].shape[0] for _ in range(2) for batch in loader]
+        assert sizes == [7, 1, 7, 1, 7, 1]
+        workers = tuple(loader.loader._iterator._workers) if num_workers else ()
+    finally:
+        loader.close()
+    assert all(not worker.is_alive() for worker in workers)
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+def test_uniform_group_combines_numpy_and_lmdb_systems(
+    uniform_sources: tuple[list[str], list[np.ndarray]],
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    backend: str,
+) -> None:
+    from deepmd.pt.utils import (
+        env,
+    )
+
+    monkeypatch.setattr(env, "NUM_WORKERS", 0)
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    paths, _ = uniform_sources
+    numpy_path = tmp_path / "numpy"
+    set_path = numpy_path / "set.000"
+    set_path.mkdir(parents=True)
+    np.savetxt(numpy_path / "type.raw", [0], fmt="%d")
+    np.save(set_path / "coord.npy", np.zeros((2, 3), dtype=np.float64))
+    np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (2, 1)))
+    np.save(set_path / "energy.npy", np.full((2, 1), 5.0, dtype=np.float64))
+    config = {
+        "systems": [str(numpy_path), *paths[:2]],
+        "batch_size": [3, 2, 2],
+        "auto_prob": "prob_uniform",
+    }
+    requirements = [
+        DataRequirementItem("energy", 1, atomic=False, must=True, high_prec=True)
+    ]
+    if backend == "pt":
+        data = create_dataset(config, ["O", "H"])
+        data.add_data_requirement(requirements)
+        owner = create_dataloader(data, config, seed=4)
+        batches = chain.from_iterable(repeat(owner))
+    else:
+        data = get_data(
+            config,
+            0,
+            ["O", "H"],
+            None,
+            seed=4,
+            lmdb_factory=lambda path, size, auto_prob: LmdbDataSystem(
+                path, ["O", "H"], size, auto_prob, seed=4, num_workers=0
+            ),
+        )
+        data.add_data_requirements(requirements)
+        owner = data
+        batches = (data.get_batch() for _ in repeat(None))
+    counts = np.zeros(6, dtype=np.int64)
+    try:
+        assert owner.children[0].uniform_sampling_mass == 0.5
+        for batch in islice(batches, 2000):
+            labels = np.asarray(batch["energy"]).reshape(-1).astype(int)
+            counts += np.bincount(labels, minlength=6)
+    finally:
+        owner.close()
+    np.testing.assert_allclose(counts[[0, 1, 2, 3, 5]] / counts.sum(), 0.2, atol=0.025)
+    assert counts[4] == 0
+
+
+@pytest.mark.parametrize("grouped", [False, True])
+def test_multitask_defaults_match_native_pt_sizes(
+    uniform_sources: tuple[list[str], list[np.ndarray]], grouped: bool
+) -> None:
+    paths, _ = uniform_sources
+    first = {"systems": paths[0], "batch_size": 1}
+    configs = {
+        "first": [{**first, "prob": 0.0}, {**first, "prob": 1.0}] if grouped else first,
+        "second": {
+            "systems": paths[1:],
+            "batch_size": "filter:12",
+            "auto_prob": "prob_uniform",
+        },
+    }
+    expected = [1 / 3, 2 / 3] if grouped else [0.2, 0.8]
+
+    def unused_epoch_length(key: str) -> int:
+        raise AssertionError(
+            f"Default model probabilities must not read epoch length: {key}"
+        )
+
+    for backend in ("pt", "pt_expt"):
+        datasets = {
+            key: (
+                create_dataset(config, ["O", "H"])
+                if backend == "pt"
+                else get_data(
+                    config,
+                    0,
+                    ["O", "H"],
+                    None,
+                    lmdb_factory=lambda path, size, auto_prob: LmdbDataSystem(
+                        path, ["O", "H"], size, auto_prob, num_workers=0
+                    ),
+                )
+            )
+            for key, config in configs.items()
+        }
+        try:
+            schedule = resolve_step_schedule(
+                {"numb_steps": 5},
+                multi_task=True,
+                model_keys=list(configs),
+                training_data=datasets,
+                epoch_length=unused_epoch_length,
+            )
+            np.testing.assert_allclose(schedule.model_prob, expected)
+        finally:
+            for data in datasets.values():
+                data.close()
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+def test_uniform_distributed_sources_use_the_queued_global_step(
+    uniform_sources: tuple[list[str], list[np.ndarray]],
+    monkeypatch: pytest.MonkeyPatch,
+    backend: str,
+) -> None:
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    paths, _ = uniform_sources
+    config = {"systems": paths, "batch_size": "mix:30", "auto_prob": "prob_uniform"}
+    requirements = [
+        DataRequirementItem("energy", 1, atomic=False, must=True, high_prec=True)
+    ]
+    owners, iterators = [], []
+    for rank in range(3):
+        if backend == "pt":
+            data = create_dataset(config, ["O", "H"])
+            data.add_data_requirement(requirements)
+            owner = create_dataloader(data, config, seed=12, rank=rank, world_size=3)
+            iterator = chain.from_iterable(repeat(owner))
+        else:
+            data = get_data(
+                config,
+                0,
+                ["O", "H"],
+                None,
+                seed=12,
+                lmdb_factory=lambda path, size, auto_prob: LmdbDataSystem(
+                    path,
+                    ["O", "H"],
+                    size,
+                    auto_prob,
+                    seed=12,
+                    rank=rank,
+                    world_size=3,
+                    num_workers=0,
+                ),
+            )
+            data.add_data_requirements(requirements)
+            data.use_ragged_batches(True)
+            owner = data
+            iterator = (source.get_batch() for source in repeat(data))
+        owners.append(owner)
+        iterators.append(iterator)
+    counts = np.zeros(5, dtype=np.int64)
+    try:
+        for _ in range(1000):
+            masses = [
+                [child.uniform_sampling_mass for child in owner.children]
+                for owner in owners
+            ]
+            np.testing.assert_allclose(masses, np.tile(masses[0], (3, 1)))
+            batches = [next(iterator) for iterator in iterators]
+            labels = [
+                np.asarray(batch["energy"]).reshape(-1).astype(int) for batch in batches
+            ]
+            sources = [
+                0 if values[0] == 0 else 2 if values[0] == 4 else 1 for values in labels
+            ]
+            assert sources[0] == sources[1] == sources[2]
+            mean_frames = np.mean([len(values) for values in labels])
+            assert masses[0][sources[0]] * mean_frames == pytest.approx(
+                [1, 3, 1][sources[0]]
+            )
+            counts += np.bincount(np.concatenate(labels), minlength=5)
+    finally:
+        for owner in owners:
+            owner.close()
+    np.testing.assert_allclose(counts / counts.sum(), 0.2, atol=0.035)
diff --git a/source/tests/pt/test_neighbor_backend_missing.py b/source/tests/pt/test_neighbor_backend_missing.py
index dd608f5e51..3948f0ac3d 100644
--- a/source/tests/pt/test_neighbor_backend_missing.py
+++ b/source/tests/pt/test_neighbor_backend_missing.py
@@ -1,48 +1,78 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""A missing neighbor-list backend must say what to install.
-
-The DPA-4 / SeZM descriptor needs ``vesin`` or ``nvalchemiops`` at runtime.
-Both ship with the ``torch`` extra, so a normal install has one, but an
-install that skips that extra reaches a model build and then fails with a
-message naming the *imports* rather than the distributions that provide
-them. These tests pin that the message names installable packages.
-"""
+"""Native fallback retains every cutoff edge without optional search libraries."""
 
 import unittest
+import warnings
+from unittest import (
+    mock,
+)
 
 import torch
 
-from deepmd.pt.model.model import (
-    sezm_model,
+from deepmd.pt_expt.utils import (
+    auto_neighbor_list,
 )
 
 
 class TestNeighborBackendMissing(unittest.TestCase):
-    def _message_without_backends(self) -> str:
-        """The error raised when neither backend is importable."""
-        nv = sezm_model.is_nv_available
-        vesin = sezm_model.is_vesin_torch_available
-        sezm_model.is_nv_available = lambda: False
-        sezm_model.is_vesin_torch_available = lambda: False
-        try:
-            with self.assertRaises(RuntimeError) as caught:
-                sezm_model._select_neighbor_builder(1, torch.device("cpu"))
-            return str(caught.exception)
-        finally:
-            sezm_model.is_nv_available = nv
-            sezm_model.is_vesin_torch_available = vesin
-
-    def test_error_names_installable_packages(self) -> None:
-        """The message must name pip-installable distributions, not imports."""
-        message = self._message_without_backends()
-        self.assertIn("pip install", message)
-        # the distribution names, which differ from the import names
-        self.assertIn("vesin[torch]", message)
-        self.assertIn("nvalchemi-toolkit-ops", message)
-
-    def test_error_points_at_the_extra(self) -> None:
-        """A user who wants the whole descriptor should be told about it."""
-        self.assertIn("deepmd-kit[torch]", self._message_without_backends())
+    def setUp(self) -> None:
+        for name in ("is_nv_available", "is_vesin_torch_available"):
+            patcher = mock.patch.object(auto_neighbor_list, name, return_value=False)
+            patcher.start()
+            self.addCleanup(patcher.stop)
+
+    def test_dense_fallback_reports_cost_without_repeating(self) -> None:
+        coord = torch.zeros((1, 1, 3), dtype=torch.float64, device="cpu")
+        atype = torch.zeros((1, 1), dtype=torch.int64, device="cpu")
+        with warnings.catch_warnings(record=True) as recorded:
+            warnings.simplefilter("default", RuntimeWarning)
+            for _ in range(3):
+                auto_neighbor_list.AdaptiveNeighborList(method="native").build(
+                    coord, atype, None, 2.0, [0], return_mode="edges"
+                )
+        self.assertEqual(len(recorded), 1)
+        self.assertIn("O(N^2)", str(recorded[0].message))
+        self.assertIn("nvalchemi-toolkit-ops", str(recorded[0].message))
+
+    def test_native_edges_are_not_truncated_to_sel(self) -> None:
+        """Batched real atoms keep all edges and phantom atoms keep none."""
+        coord = torch.tensor(
+            [[[0.0, 0.0, 0.0], [0.8, 0.0, 0.0], [0.0, 0.9, 0.0], [0.1, 0.0, 0.0]]],
+            dtype=torch.float64,
+            device="cpu",
+        )
+        coord = coord.repeat(2, 1, 1).requires_grad_()
+        atype = torch.tensor([[0, 1, 0, -1]], device="cpu").repeat(2, 1)
+        for method in ("native", "auto"):
+            with self.subTest(method=method):
+                builder = auto_neighbor_list.make_neighbor_list(
+                    method, coord.device, return_mode="edges"
+                )
+                edges = builder.build(
+                    coord, atype, None, 2.0, [1, 0], return_mode="edges"
+                )
+                indices = edges.edge_index[:, edges.edge_mask]
+                self.assertEqual(indices.shape[1], 12)
+                self.assertTrue((atype.reshape(-1)[indices] >= 0).all())
+                self.assertTrue((indices[0] // 4 == indices[1] // 4).all())
+                grad = torch.autograd.grad(edges.edge_vec.square().sum(), coord)[0]
+                self.assertTrue(torch.isfinite(grad).all())
+                self.assertGreater(float(grad.abs().sum()), 0.0)
+
+    def test_native_edges_keep_periodic_self_images(self) -> None:
+        coord = torch.zeros((1, 1, 3), dtype=torch.float64, device="cpu")
+        atype = torch.zeros((1, 1), dtype=torch.int64, device="cpu")
+        box = torch.eye(3, dtype=coord.dtype, device="cpu").unsqueeze(0)
+        builder = auto_neighbor_list.make_neighbor_list(
+            "auto", coord.device, return_mode="edges"
+        )
+        edges = builder.build(coord, atype, box, 1.1, [0], return_mode="edges")
+        self.assertEqual(int(edges.edge_mask.sum()), 6)
+        self.assertTrue((edges.edge_index[:, edges.edge_mask] == 0).all())
+        torch.testing.assert_close(
+            edges.edge_vec[edges.edge_mask].square().sum(-1),
+            torch.ones(6, dtype=coord.dtype, device="cpu"),
+        )
 
 
 if __name__ == "__main__":
diff --git a/source/tests/pt/test_sampler.py b/source/tests/pt/test_sampler.py
index 88848c7e02..ee35acd6f8 100644
--- a/source/tests/pt/test_sampler.py
+++ b/source/tests/pt/test_sampler.py
@@ -17,6 +17,9 @@
 from deepmd.dpmodel.utils import (
     compute_total_numb_batch,
 )
+from deepmd.dpmodel.utils.training_utils import (
+    resolve_model_prob,
+)
 from deepmd.pt.utils import (
     dp_random,
 )
@@ -168,6 +171,29 @@ def test_sampler_debug_info(self) -> None:
         frame = self.my_dataset.systems[sid].__getitem__(fid)
         self.assertTrue(np.allclose(coord, frame["coord"]))
 
+    def test_grouping_preserves_default_model_probabilities(self) -> None:
+        system = self.systems[0]
+        grouped = pt_dataloader.create_dataset(
+            [
+                {"systems": [system] * 2, "batch_size": 1},
+                {"systems": [system] * 3, "batch_size": 1},
+            ],
+            self.type_map,
+        )
+        flat = pt_dataloader.create_dataset(
+            {"systems": [system] * 5, "batch_size": 1}, self.type_map
+        )
+        single = pt_dataloader.create_dataset(
+            {"systems": system, "batch_size": 1}, self.type_map
+        )
+        probabilities = [
+            resolve_model_prob(
+                ["large", "small"], None, {"large": data, "small": single}
+            )
+            for data in (flat, grouped)
+        ]
+        np.testing.assert_allclose(probabilities, [[5 / 6, 1 / 6]] * 2)
+
     def test_auto_prob_uniform(self) -> None:
         auto_prob_style = "prob_uniform"
         sampler = pt_dataloader.get_weighted_sampler(
diff --git a/source/tests/pt/test_tabulate_gpu_size_validation.py b/source/tests/pt/test_tabulate_gpu_size_validation.py
new file mode 100644
index 0000000000..f75314e9a3
--- /dev/null
+++ b/source/tests/pt/test_tabulate_gpu_size_validation.py
@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Reject invalid tabulation widths before any PyTorch GPU kernel launch."""
+
+import unittest
+from collections.abc import (
+    Callable,
+)
+
+import torch
+
+from deepmd.pt.cxx_op import (
+    ENABLE_CUSTOMIZED_OP,
+)
+
+ERROR_MESSAGE = "last_layer_size must be between 1 and 1024"
+POSITIVE_ERROR_MESSAGE = "last_layer_size must be positive"
+
+
+@unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built")
+@unittest.skipUnless(
+    torch.cuda.is_available(), "GPU tabulation validation requires a GPU"
+)
+class TestTabulateGpuSizeValidation(unittest.TestCase):
+    """The GPU wrappers must fail on the host, not inside a launch config."""
+
+    dtype = torch.float64
+
+    def _zeros(self, *shape: int) -> torch.Tensor:
+        return torch.zeros(shape, dtype=self.dtype, device="cuda")
+
+    def _table(self, last_layer_size: int) -> torch.Tensor:
+        return self._zeros(1, max(1, 6 * last_layer_size))
+
+    def _table_info(self) -> torch.Tensor:
+        # table_info stays on the CPU for every backend.
+        return torch.tensor(
+            [0.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=self.dtype, device="cpu"
+        )
+
+    def _forward_calls(self, last_layer_size: int) -> dict[str, Callable[[], object]]:
+        table = self._table(last_layer_size)
+        table_info = self._table_info()
+        em_x = self._zeros(1, 1)
+        em_a = self._zeros(1, 1, 4)
+        em_t = self._zeros(1, 1, 1)
+        em_r = self._zeros(1, 1)
+        two_embed = self._zeros(1, max(1, last_layer_size))
+        return {
+            "se_a": lambda: torch.ops.deepmd.tabulate_fusion_se_a(
+                table, table_info, em_x, em_a, last_layer_size
+            ),
+            "se_atten": lambda: torch.ops.deepmd.tabulate_fusion_se_atten(
+                table, table_info, em_x, em_a, two_embed, last_layer_size, True
+            ),
+            "se_t": lambda: torch.ops.deepmd.tabulate_fusion_se_t(
+                table, table_info, em_x, em_t, last_layer_size
+            ),
+            "se_t_tebd": lambda: torch.ops.deepmd.tabulate_fusion_se_t_tebd(
+                table, table_info, em_x, em_t, last_layer_size
+            ),
+            "se_r": lambda: torch.ops.deepmd.tabulate_fusion_se_r(
+                table, table_info, em_r, last_layer_size
+            ),
+        }
+
+    def _assert_rejected(
+        self,
+        last_layer_size: int,
+        *,
+        include_se_t_tebd: bool = True,
+    ) -> None:
+        calls = self._forward_calls(last_layer_size)
+        if not include_se_t_tebd:
+            calls.pop("se_t_tebd")
+        for name, call in calls.items():
+            with self.subTest(op=name, last_layer_size=last_layer_size):
+                error_message = (
+                    POSITIVE_ERROR_MESSAGE if name == "se_t_tebd" else ERROR_MESSAGE
+                )
+                with self.assertRaisesRegex(RuntimeError, error_message):
+                    call()
+
+    def test_rejects_oversized_width(self) -> None:
+        """A width past the maximum block dimension must be refused."""
+        self._assert_rejected(1025, include_se_t_tebd=False)
+
+    def test_rejects_nonpositive_width(self) -> None:
+        """Zero and negative widths are invalid launch dimensions."""
+        self._assert_rejected(0)
+        self._assert_rejected(-1)
+
+    def _gradient_cases(
+        self, last_layer_size: int = 2
+    ) -> dict[
+        str,
+        tuple[
+            Callable[[], torch.Tensor],
+            tuple[torch.Tensor, ...],
+            tuple[int, ...],
+        ],
+    ]:
+        """Build valid forwards whose saved descriptors can test grad guards."""
+        table = self._table(last_layer_size)
+        table_info = self._table_info()
+        em_x = self._zeros(1, 1).requires_grad_(True)
+        em_a = self._zeros(1, 1, 4).requires_grad_(True)
+        two_embed = self._zeros(1, last_layer_size).requires_grad_(True)
+        em_t = self._zeros(1, 1, 1).requires_grad_(True)
+        em_t_tebd_x = self._zeros(1, 1).requires_grad_(True)
+        em_t_tebd = self._zeros(1, 1, 1).requires_grad_(True)
+        em_r = self._zeros(1, 1).requires_grad_(True)
+        return {
+            "se_a": (
+                lambda: torch.ops.deepmd.tabulate_fusion_se_a(
+                    table, table_info, em_x, em_a, last_layer_size
+                )[0],
+                (em_x, em_a),
+                (1, 4, last_layer_size),
+            ),
+            "se_atten": (
+                lambda: torch.ops.deepmd.tabulate_fusion_se_atten(
+                    table,
+                    table_info,
+                    em_x,
+                    em_a,
+                    two_embed,
+                    last_layer_size,
+                    True,
+                )[0],
+                (em_x, em_a, two_embed),
+                (1, 4, last_layer_size),
+            ),
+            "se_t": (
+                lambda: torch.ops.deepmd.tabulate_fusion_se_t(
+                    table, table_info, em_x, em_t, last_layer_size
+                )[0],
+                (em_x, em_t),
+                (1, last_layer_size),
+            ),
+            "se_t_tebd": (
+                lambda: torch.ops.deepmd.tabulate_fusion_se_t_tebd(
+                    table,
+                    table_info,
+                    em_t_tebd_x,
+                    em_t_tebd,
+                    last_layer_size,
+                )[0],
+                (em_t_tebd_x,),
+                (1, 1, 1, last_layer_size),
+            ),
+            "se_r": (
+                lambda: torch.ops.deepmd.tabulate_fusion_se_r(
+                    table, table_info, em_r, last_layer_size
+                )[0],
+                (em_r,),
+                (1, 1, last_layer_size),
+            ),
+        }
+
+    def _oversized_saved_descriptor(
+        self, descriptor_shape: tuple[int, ...]
+    ) -> Callable[[torch.Tensor], torch.Tensor]:
+        """Replace only a saved descriptor with an oversized-width sentinel."""
+
+        def pack(tensor: torch.Tensor) -> torch.Tensor:
+            if tuple(tensor.shape) == descriptor_shape:
+                return self._zeros(*descriptor_shape[:-1], 1025)
+            return tensor
+
+        return pack
+
+    @staticmethod
+    def _unpack_saved_tensor(tensor: torch.Tensor) -> torch.Tensor:
+        """Return saved tensors unchanged when their autograd node reloads."""
+        return tensor
+
+    def test_first_gradient_wrappers_reject_oversized_width(self) -> None:
+        """Every first-gradient wrapper validates its saved descriptor width."""
+        for name, (forward, inputs, descriptor_shape) in self._gradient_cases().items():
+            if name == "se_t_tebd":
+                continue
+            with self.subTest(op=name):
+                pack = self._oversized_saved_descriptor(descriptor_shape)
+                with torch.autograd.graph.saved_tensors_hooks(
+                    pack, self._unpack_saved_tensor
+                ):
+                    descriptor = forward()
+                with self.assertRaisesRegex(RuntimeError, ERROR_MESSAGE):
+                    torch.autograd.grad(descriptor.sum(), inputs)
+
+    def test_second_gradient_wrappers_reject_oversized_width(self) -> None:
+        """Every grad-grad wrapper validates the descriptor before launching."""
+        for name, (forward, inputs, descriptor_shape) in self._gradient_cases().items():
+            if name == "se_t_tebd":
+                continue
+            with self.subTest(op=name):
+                descriptor = forward()
+                pack = self._oversized_saved_descriptor(descriptor_shape)
+                with torch.autograd.graph.saved_tensors_hooks(
+                    pack, self._unpack_saved_tensor
+                ):
+                    first_gradients = torch.autograd.grad(
+                        descriptor.sum(), inputs, create_graph=True
+                    )
+                differentiable_sum = sum(
+                    gradient.sum()
+                    for gradient in first_gradients
+                    if gradient.requires_grad
+                )
+                with self.assertRaisesRegex(RuntimeError, ERROR_MESSAGE):
+                    torch.autograd.grad(differentiable_sum, inputs)
+
+    def test_se_t_tebd_accepts_oversized_width(self) -> None:
+        """SE-T-TEBD uses fixed-size GPU blocks for forward and gradients."""
+        last_layer_size = 1025
+        forward, inputs, _ = self._gradient_cases(last_layer_size)["se_t_tebd"]
+        descriptor = forward()
+        self.assertEqual(descriptor.shape[-1], last_layer_size)
+
+        first_gradients = torch.autograd.grad(
+            descriptor.sum(), inputs, create_graph=True
+        )
+        differentiable_sum = sum(
+            gradient.sum() for gradient in first_gradients if gradient.requires_grad
+        )
+        second_gradients = torch.autograd.grad(differentiable_sum, inputs)
+        self.assertEqual(len(second_gradients), len(inputs))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py
index 85a8a1a749..052d0166c6 100644
--- a/source/tests/pt/test_training.py
+++ b/source/tests/pt/test_training.py
@@ -1042,6 +1042,37 @@ def tearDown(self) -> None:
         os.chdir(self._cwd)
         self._tmpdir.cleanup()
 
+    @TRAINING_TEST_TIMEOUT
+    @patch("deepmd.pt_expt.train.validation.FullValidator.evaluate_all_systems")
+    def test_full_validation_rejects_prior_nonfinite_gradient(
+        self, mocked_eval
+    ) -> None:
+        mocked_eval.return_value = {"mae_e_per_atom": 1.0}
+        for ema in (False, True):
+            with self.subTest(ema=ema):
+                config = deepcopy(self.config)
+                config["training"]["enable_ema"] = True
+                config["validating"]["full_validation"] = not ema
+                config["validating"]["ema_full_validation"] = ema
+                trainer = get_trainer(config)
+                validator = (
+                    trainer.ema_full_validator if ema else trainer.full_validator
+                )
+                previous_state = deepcopy(validator.state_store)
+                trainer.nonfinite_grad_guard.update(
+                    torch.tensor(float("nan"), device="cpu")
+                )
+                trainer.nonfinite_grad_guard.update(torch.tensor(1.0, device="cpu"))
+                with (
+                    patch("torch.save", wraps=torch.save) as serialize,
+                    self.assertRaisesRegex(RuntimeError, "Non-finite gradient norm"),
+                ):
+                    trainer.run()
+                serialize.assert_not_called()
+                self.assertEqual(validator.state_store, previous_state)
+                self.assertEqual(validator.topk_records, [])
+                self.assertEqual(list(Path(".").glob("*.pt")), [])
+
     @TRAINING_TEST_TIMEOUT
     @patch("deepmd.pt_expt.train.validation.FullValidator.evaluate_all_systems")
     def test_full_validation_rotates_best_checkpoint(self, mocked_eval) -> None:
@@ -1267,15 +1298,11 @@ def get_data(
 
 class TestEMATraining(unittest.TestCase):
     def setUp(self) -> None:
-        import deepmd.pt.train.training as training_module
         import deepmd.pt.utils.env as env_module
 
-        self._num_workers_state = (
-            (env_module, env_module.NUM_WORKERS),
-            (training_module, training_module.NUM_WORKERS),
-        )
+        self._num_workers_state = env_module.NUM_WORKERS
+        self.addCleanup(setattr, env_module, "NUM_WORKERS", self._num_workers_state)
         env_module.NUM_WORKERS = 0
-        training_module.NUM_WORKERS = 0
         self._cwd = os.getcwd()
         self._tmpdir = tempfile.TemporaryDirectory()
         os.chdir(self._tmpdir.name)
@@ -1306,8 +1333,6 @@ def setUp(self) -> None:
         }
 
     def tearDown(self) -> None:
-        for module, num_workers in self._num_workers_state:
-            module.NUM_WORKERS = num_workers
         os.chdir(self._cwd)
         self._tmpdir.cleanup()
 
diff --git a/source/tests/pt/test_update_sel.py b/source/tests/pt/test_update_sel.py
index c87c3330e9..e34672c150 100644
--- a/source/tests/pt/test_update_sel.py
+++ b/source/tests/pt/test_update_sel.py
@@ -143,33 +143,6 @@ def test_update_sel_atten_int(self, sel_mock) -> None:
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat")
-    def test_update_sel_atten_list(self, sel_mock) -> None:
-        sel_mock.return_value = self.mock_min_nbor_dist, [25]
-
-        jdata = {
-            "model": {
-                "descriptor": {
-                    "type": "se_atten",
-                    "sel": 30,
-                    "rcut": 6,
-                }
-            },
-            "training": {"training_data": {}},
-        }
-        expected_out = {
-            "model": {
-                "descriptor": {
-                    "type": "se_atten",
-                    "sel": 30,
-                    "rcut": 6,
-                }
-            },
-            "training": {"training_data": {}},
-        }
-        jdata = update_sel(jdata)
-        self.assertEqual(jdata, expected_out)
-
     @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat")
     def test_update_sel_dpa2_auto(self, sel_mock) -> None:
         sel_mock.return_value = self.mock_min_nbor_dist, [25]
diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py
index 767b9967e4..d90087c055 100644
--- a/source/tests/pt/test_validation.py
+++ b/source/tests/pt/test_validation.py
@@ -268,6 +268,35 @@ def _make_spin_task_config() -> dict:
 
 
 class TestValidationHelpers(unittest.TestCase):
+    @staticmethod
+    def _run_best_candidate(
+        validator: FullValidator,
+        *,
+        display_step: int,
+        selected_metric_value: float,
+        directory_checkpoint: bool = False,
+    ) -> str:
+        """Exercise selection, writing and reconciliation through the public flow."""
+
+        def save_checkpoint(path: Path, **kwargs) -> None:
+            if directory_checkpoint:
+                path.mkdir()
+            else:
+                path.touch()
+
+        with patch.object(
+            validator,
+            "evaluate_all_systems",
+            return_value={validator.metric_key: selected_metric_value},
+        ):
+            result = validator.run(
+                step_id=display_step,
+                display_step=display_step,
+                lr=1e-3,
+                save_checkpoint=save_checkpoint,
+            )
+        return result.saved_best_path
+
     def test_resolve_full_validation_start_step(self) -> None:
         self.assertEqual(resolve_full_validation_start_step(0, 2000000), 0)
         self.assertEqual(resolve_full_validation_start_step(0.1, 2000000), 200000)
@@ -349,26 +378,23 @@ def test_full_validator_rotates_best_checkpoint(self) -> None:
                     rank=0,
                     restart_training=False,
                 )
-                new_best_path = validator._update_best_state(
+                self._run_best_candidate(
+                    validator,
                     display_step=1,
                     selected_metric_value=2.0,
                 )
-                Path(new_best_path).touch()
-                validator._reconcile_best_checkpoints()
 
-                new_best_path = validator._update_best_state(
+                self._run_best_candidate(
+                    validator,
                     display_step=2,
                     selected_metric_value=1.0,
                 )
-                Path(new_best_path).touch()
-                validator._reconcile_best_checkpoints()
 
-                new_best_path = validator._update_best_state(
+                new_best_path = self._run_best_candidate(
+                    validator,
                     display_step=3,
                     selected_metric_value=1.5,
                 )
-                Path(new_best_path).touch()
-                validator._reconcile_best_checkpoints()
             finally:
                 os.chdir(old_cwd)
 
@@ -453,7 +479,8 @@ def test_full_validator_writes_best_into_custom_checkpoint_dir(self) -> None:
                 )
                 # The directory is created recursively at construction time.
                 self.assertTrue(best_dir.is_dir())
-                new_best_path = validator._update_best_state(
+                new_best_path = self._run_best_candidate(
+                    validator,
                     display_step=1,
                     selected_metric_value=2.0,
                 )
@@ -486,26 +513,26 @@ def test_full_validator_reconciles_directory_checkpoints(self) -> None:
                     restart_training=False,
                     best_checkpoint_suffix=".jax",
                 )
-                new_best_path = validator._update_best_state(
+                self._run_best_candidate(
+                    validator,
+                    directory_checkpoint=True,
                     display_step=1,
                     selected_metric_value=2.0,
                 )
-                Path(new_best_path).mkdir()
-                validator._reconcile_best_checkpoints()
 
-                new_best_path = validator._update_best_state(
+                self._run_best_candidate(
+                    validator,
+                    directory_checkpoint=True,
                     display_step=2,
                     selected_metric_value=1.0,
                 )
-                Path(new_best_path).mkdir()
-                validator._reconcile_best_checkpoints()
 
-                new_best_path = validator._update_best_state(
+                new_best_path = self._run_best_candidate(
+                    validator,
+                    directory_checkpoint=True,
                     display_step=3,
                     selected_metric_value=1.5,
                 )
-                Path(new_best_path).mkdir()
-                validator._reconcile_best_checkpoints()
             finally:
                 os.chdir(old_cwd)
 
diff --git a/source/tests/pt/test_wrapper.py b/source/tests/pt/test_wrapper.py
index 9fd64bdb06..5d74ff7b1b 100644
--- a/source/tests/pt/test_wrapper.py
+++ b/source/tests/pt/test_wrapper.py
@@ -21,6 +21,7 @@ def __init__(self, *, fail_forward: bool = False) -> None:
         self.scale = torch.nn.Parameter(torch.ones((), device="cpu"))
         self.fail_forward = fail_forward
         self.last_requires_grad: tuple[bool, ...] | None = None
+        self.last_box: torch.Tensor | None = None
 
     def forward(
         self,
@@ -32,7 +33,8 @@ def forward(
         aparam: torch.Tensor | None = None,
         charge_spin: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
-        del atype, box, do_atomic_virial, fparam, aparam, charge_spin
+        del atype, do_atomic_virial, fparam, aparam, charge_spin
+        self.last_box = box
         self.last_requires_grad = tuple(
             param.requires_grad for param in self.parameters()
         )
@@ -131,6 +133,15 @@ def test_skip_loss_keeps_training_gradients(self) -> None:
         self.assertIsNotNone(model.linear.weight.grad)
         self.assertIsNotNone(model.scale.grad)
 
+    def test_positional_box_keeps_legacy_argument_order(self) -> None:
+        model = _LinearToyModel()
+        wrapper = ModelWrapper(model)
+        box = torch.eye(3, device="cpu").view(1, 3, 3).expand(2, -1, -1)
+
+        wrapper(self.coord, self.atype, None, box)
+
+        self.assertIs(model.last_box, box)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py
index 2f1f81afe5..4f801e328c 100644
--- a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py
+++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py
@@ -34,7 +34,14 @@
 import copy
 import dataclasses
 import os
+import tempfile
 import unittest
+from pathlib import (
+    Path,
+)
+from unittest import (
+    mock,
+)
 
 import torch
 from torch.fx.experimental.proxy_tensor import (
@@ -1138,6 +1145,150 @@ def test_parity_single_layer_residual(self) -> None:
             )
         )
 
+    def test_fused_network_matches_the_layer_kernels_bitwise(self) -> None:
+        """The fused network reproduces the per-layer kernels bit for bit.
+
+        ``graph_fitting_energy_gradient`` evaluates a run with the fused
+        network wherever that serves it and with the per-layer kernels
+        elsewhere, while the separate forward and backward operators always
+        take the per-layer kernels. The runs, sized by the multiprocessor
+        count, span the shapes the dispatch selects by run size: a cluster or
+        latency shape up to one panel per multiprocessor, the throughput shape
+        of the widths that have one beyond, at a width of 96 the 32-node panels
+        on a run that one wave of either throughput shape holds and the 64-node
+        panels on runs that they pass in fewer waves or in two waves or more,
+        and the per-layer kernels wherever the fused network declines a run.
+        The tiled calls end in a short tail, one of them after a first run of
+        exactly two waves of 64-node panels.
+        Four networks cover the input paths and the depth: an input narrower
+        than the width and not a multiple of four, an aligned input wider than
+        the width, whose layer-0 backward contracts in several passes, a single
+        square layer without biases, whose backward the residual seeds, and
+        the eight layers the fused network holds at most. A network whose runs
+        the fused network never serves, as on a device or build without a
+        kernel image of compute capability 7.0 or later, is reported as
+        skipped, since its comparison covers the per-layer kernels alone. The
+        route of a run is that of the rules of the dispatch
+        (``graph_fitting_route_rule``), which the call shares.
+        """
+        from deepmd.pt_expt.kernels.graph_fitting import (
+            ensure_registered,
+        )
+
+        ensure_registered()
+        sms = torch.cuda.get_device_properties(self.device).multi_processor_count
+        runs = (
+            (37, 0),
+            (16 * sms + 3, 0),
+            (32 * sms - 1, 0),
+            (48 * sms + 17, 0),
+            (96 * sms + 17, 0),
+            (96 * sms + 17, 32 * sms - 1),
+            (256 * sms + 37, 0),
+            (256 * sms + 37, 256 * sms),
+        )
+        gen = torch.Generator(device=self.device).manual_seed(11)
+        for width in (64, 96, 192, 256, 384):
+            networks = (
+                (70, 3, True),
+                (2 * width + 8, 3, True),
+                (width, 1, False),
+                (70, 8, True),
+            )
+            for act in ("tanh", "silu"):
+                for din, depth, biased in networks:
+                    self._assert_fused_matches_layers(
+                        width, act, din, depth, biased, runs, gen
+                    )
+
+    def _assert_fused_matches_layers(
+        self,
+        width: int,
+        act: str,
+        din: int,
+        depth: int,
+        biased: bool,
+        runs: tuple[tuple[int, int], ...],
+        gen: torch.Generator,
+    ) -> None:
+        """Compare the fused and per-layer kernels on one network."""
+        from deepmd.pt_expt.kernels.fitting_routes import (
+            PER_LAYER,
+        )
+        from deepmd.pt_expt.kernels.graph_fitting import (
+            fitting_operator_arguments,
+        )
+
+        fit = self._build(
+            resnet_dt=False, act=act, dim_descrpt=din, neuron=[width] * depth
+        )
+        network = fitting_operator_arguments(fit)
+        biases = (
+            network.biases
+            if biased
+            else [layer_bias.new_empty(0) for layer_bias in network.biases]
+        )
+        head_bias = network.head_bias if biased else network.head_bias.new_empty(0)
+        bias = fit.bias_atom_e[:, 0].contiguous()
+        fused_runs = 0
+        for nodes, tile in runs:
+            x = torch.randn(nodes, din, generator=gen, device=self.device)
+            atype = torch.randint(0, 2, (nodes,), generator=gen, device=self.device)
+            seed = torch.rand(
+                nodes, generator=gen, device=self.device, dtype=torch.float64
+            )
+            with torch.no_grad():
+                layer_energy, saved = torch.ops.deepmd.graph_fitting(
+                    x,
+                    atype,
+                    network.weights,
+                    biases,
+                    network.residuals,
+                    network.head_weight,
+                    head_bias,
+                    bias,
+                    network.activation,
+                )
+                layer_gradient = torch.ops.deepmd.graph_fitting_backward(
+                    seed[:, None],
+                    saved,
+                    network.weights,
+                    biases,
+                    network.residuals,
+                    network.head_weight,
+                    network.activation,
+                )
+                gradient = x.clone()
+                energy = torch.ops.deepmd.graph_fitting_energy_gradient(
+                    gradient,
+                    atype,
+                    network.weights,
+                    biases,
+                    network.residuals,
+                    network.head_weight,
+                    head_bias,
+                    bias,
+                    network.activation,
+                    seed,
+                    tile,
+                )
+            # The operator serves every run of a call with the fused network
+            # where its rules route a run of the call's run size there.
+            run = min(tile, nodes) if tile > 0 else nodes
+            fused_runs += (
+                torch.ops.deepmd.graph_fitting_route_rule(network.weights, run)
+                != PER_LAYER
+            )
+            with self.subTest(width=width, act=act, din=din, nodes=nodes, tile=tile):
+                self.assertTrue(torch.equal(energy, layer_energy))
+                self.assertTrue(torch.equal(gradient, layer_gradient))
+        with self.subTest(width=width, act=act, din=din, depth=depth):
+            if fused_runs == 0:
+                self.skipTest(
+                    "the fused network serves no run of this network on this "
+                    "device and build"
+                )
+
     def test_timestep_falls_back(self) -> None:
         from deepmd.pt_expt.kernels.graph_fitting import (
             fitting_eligible,
@@ -1711,6 +1862,79 @@ def test_compact_canonical_model_trace(self) -> None:
         for key in actual:
             torch.testing.assert_close(actual[key], reference[key])
 
+    def test_compact_canonical_deep_eval(self) -> None:
+        """A compiled canonical archive reaches the public evaluator unchanged."""
+        from deepmd.infer import (
+            DeepPot,
+        )
+        from deepmd.pt_expt.model import (
+            EnergyModel,
+        )
+        from deepmd.pt_expt.utils.serialization import (
+            deserialize_to_file,
+        )
+
+        descriptor = _build_compressed_dpa1(
+            self.device, [16, 32, 64], act="silu", axis_neuron=16
+        )
+        model = (
+            EnergyModel(
+                descriptor,
+                self._build_fitting(descriptor.get_dim_out()),
+                type_map=["A", "B"],
+            )
+            .to(self.device)
+            .eval()
+        )
+        with (
+            _CudaLevel("2"),
+            torch.device(self.device),
+            tempfile.TemporaryDirectory() as directory,
+        ):
+            reference = model(self.coord, self.atype, self.box, do_atomic_virial=True)
+            path = str(Path(directory) / "canonical.pt2")
+            deserialize_to_file(
+                path,
+                {"model": model.serialize()},
+                lower_kind="auto",
+                do_atomic_virial=True,
+                with_comm=False,
+            )
+            dp = DeepPot(path, auto_batch_size=False, neighbor_graph_method="dense")
+            self.assertEqual(
+                dp.deep_eval.metadata["lower_input_kind"], "dpa1_canonical"
+            )
+            with mock.patch.object(
+                dp.deep_eval, "_pt2_runner", wraps=dp.deep_eval._pt2_runner
+            ) as run:
+                actual = dp.eval(
+                    self.coord.cpu().numpy(),
+                    self.box.cpu().numpy(),
+                    self.atype[0].cpu().numpy(),
+                    atomic=True,
+                )
+            arguments = run.call_args.args
+            self.assertEqual(len(arguments), 8)
+            # A small graph passes its uint32 source words as their bytes.
+            self.assertEqual(arguments[3].dtype, torch.uint8)
+            self.assertEqual(arguments[3].shape[-1], 4)
+            self.assertEqual(arguments[7].dtype, torch.int64)
+            self.assertEqual(arguments[4].dtype, torch.float32)
+            self.assertGreater(float(reference["force"].abs().max()), 1e-8)
+            for value, key in zip(
+                actual,
+                ("energy", "force", "virial", "atom_energy", "atom_virial"),
+                strict=True,
+            ):
+                torch.testing.assert_close(
+                    torch.as_tensor(
+                        value, dtype=reference[key].dtype, device=self.device
+                    ),
+                    reference[key],
+                    atol=1e-4,
+                    rtol=1e-4,
+                )
+
     def test_compact_canonical_torch_export_contract(self) -> None:
         """torch.export records the fixed eight-tensor deployment ABI."""
         from deepmd.pt_expt.model import (
@@ -1741,7 +1965,8 @@ def test_compact_canonical_torch_export_contract(self) -> None:
 
         self.assertEqual(metadata["lower_input_kind"], "dpa1_canonical")
         self.assertEqual(metadata["graph_edge_dtype"], "float32")
-        self.assertEqual(metadata["canonical_index_dtype"], "uint32")
+        self.assertEqual(metadata["canonical_index_dtype"], "auto")
+        self.assertEqual(metadata["canonical_index_storage"], "byte_words")
         self.assertNotIn("graph_index_dtype", metadata)
         self.assertEqual(
             output_keys,
diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py
index 4eef728456..eaa3b539a6 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4c.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4c.py
@@ -265,6 +265,22 @@ def test_serialization_preserves_parameters(self) -> None:
                 original_parameters[name],
             )
 
+    def test_state_dict_of_version_3_2_loads(self) -> None:
+        """A checkpoint holding the gather tables of DeePMD-kit 3.2 loads strictly."""
+        state = self.descriptor.state_dict()
+        width = 32
+        state["angular_channel_index"] = torch.arange(width)
+        state["angular_harmonic_index"] = torch.arange(width)
+        restored = self.build()
+        restored.load_state_dict(state)
+        torch.testing.assert_close(
+            self._evaluate(restored, self.coord),
+            self._evaluate(self.descriptor, self.coord),
+        )
+        state["unknown"] = torch.zeros(1)
+        with pytest.raises(RuntimeError, match="unknown"):
+            restored.load_state_dict(state)
+
     def test_deserialization_preserves_version_one_raw_radial_basis(self) -> None:
         data = self.descriptor.serialize()
         radial_basis = data["radial_basis"]
diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py b/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py
index 38c885b6f6..653991ca9e 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py
@@ -9,6 +9,7 @@
 different compiled body or a different addressing mode.
 """
 
+import math
 from typing import (
     Any,
 )
@@ -27,8 +28,12 @@
     DescrptDPA4C,
 )
 from deepmd.pt_expt.kernels.dpa4c.graph_compress import (
+    _portable_inputs,
+    _radial_knots,
     _reference_descriptor,
     build_compression_artifacts,
+    build_low_rank_table,
+    build_radial_table,
     contact_radius_input,
     descriptor_profile,
     ensure_registered,
@@ -173,6 +178,8 @@ def _arguments(descriptor: DescrptDPA4C, graph, atype: torch.Tensor) -> tuple:
         graph.destination_row_ptr,
         atype,
         artifacts["data"],
+        artifacts["basis_table"],
+        artifacts["basis_matrix"],
         artifacts["pair_film"],
         artifacts["pair_mixing"],
         artifacts["type_embedding"],
@@ -208,8 +215,13 @@ def _descriptor_backward(
 
 
 def _spin_free(arguments: tuple) -> tuple:
-    """Drop the native spin block, which the CPU kernel does not implement."""
-    return (*arguments[:15], *arguments[18:])
+    """Drop the low-rank table and the native spin block for the CPU reference.
+
+    The reference reads the radial table itself, and the CPU kernel implements
+    no spin.
+    """
+    portable = _portable_inputs(arguments)
+    return (*portable[:15], *portable[18:])
 
 
 def _check_parity(descriptor: DescrptDPA4C, canonical: bool) -> None:
@@ -660,3 +672,139 @@ def test_empty_graph_produces_the_isolated_atom_descriptor() -> None:
     assert torch.isfinite(state).all()
     torch.testing.assert_close(state[:, 0], torch.full_like(state[:, 0], 0.5))
     torch.testing.assert_close(state[:, 1], torch.full_like(state[:, 1], 0.5))
+
+
+def _restrict_scalar_rank(descriptor: DescrptDPA4C, rank: int) -> None:
+    """Confine the scalar-only radial channels to a span of ``rank`` functions.
+
+    The radial map is the bias-free output layer applied to a shared hidden
+    state, so replacing the output columns from ``C1`` on by combinations of
+    ``rank`` of them confines those channels to the span of ``rank`` functions
+    at every distance.
+    """
+    first = descriptor_profile(descriptor.channels, descriptor.lmax).degree_channels[1]
+    weight = descriptor.radial_embedding.layers[-1].w
+    generator = torch.Generator(device=weight.device).manual_seed(31)
+    mixing = torch.randn(
+        rank,
+        descriptor.channels - first,
+        generator=generator,
+        device=weight.device,
+        dtype=weight.dtype,
+    ) / math.sqrt(rank)
+    with torch.no_grad():
+        weight[:, first:] = weight[:, first : first + rank] @ mixing
+
+
+def _channel_coefficients(table: torch.Tensor, width: int) -> torch.Tensor:
+    """Return the quintic coefficients of a table with shape ``(S, width, 6)``.
+
+    A table row holds the leading four coefficients of every channel, then the
+    trailing two.
+    """
+    quartets = table[:, : 4 * width].reshape(table.shape[0], width, 4)
+    pairs = table[:, 4 * width :].reshape(table.shape[0], width, 2)
+    return torch.cat([quartets, pairs], dim=-1)
+
+
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes"),
+    [(64, 2, 0), (64, 3, 4), (128, 3, 4)],
+)
+def test_low_rank_table_reconstructs_the_radial_table(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+) -> None:
+    """The low-rank basis and its expansion reproduce the scalar-only channels.
+
+    The quintic coefficients are linear in the knot data, so the expansion of
+    the basis coefficients reproduces the coefficients of every scalar-only
+    channel of the radial table up to the rounding of their fp32 storage, and
+    the degree-one channels and the mode profiles are tabulated unchanged.
+    """
+    descriptor = _build_descriptor(channels, lmax, radial_modes)
+    _restrict_scalar_rank(descriptor, rank=8)
+    basis_table, basis_matrix = build_low_rank_table(descriptor)
+    direct, _info = build_radial_table(descriptor)
+    first = descriptor_profile(channels, lmax).degree_channels[1]
+    basis = basis_matrix.shape[0]
+    assert basis_matrix.shape == (32, channels - first)
+    direct_coefficients = _channel_coefficients(
+        direct.double(), channels + radial_modes
+    )
+    basis_coefficients = _channel_coefficients(
+        basis_table.double(), first + basis + radial_modes
+    )
+    assert torch.equal(basis_coefficients[:, :first], direct_coefficients[:, :first])
+    assert torch.equal(
+        basis_coefficients[:, first + basis :], direct_coefficients[:, channels:]
+    )
+    expanded = torch.einsum(
+        "skj,kc->scj",
+        basis_coefficients[:, first : first + basis],
+        basis_matrix.double(),
+    )
+    scalar = direct_coefficients[:, first:channels]
+    scale = scalar.abs().amax(dim=0, keepdim=True).clamp_min(1e-30)
+    torch.testing.assert_close(expanded / scale, scalar / scale, atol=1e-5, rtol=0)
+
+
+@pytest.mark.parametrize("channels", [64, 128])
+def test_full_rank_radial_map_keeps_the_direct_table(channels: int) -> None:
+    """A radial map whose scalar-only channels exceed the basis keeps none.
+
+    The knot data of the scalar-only channels is replaced by independent
+    random samples, which span as many functions as there are such channels,
+    more than the basis holds, so the compression leaves both low-rank
+    artifacts empty and the single-type scans read the radial table.
+    """
+    descriptor = _build_descriptor(channels)
+    first = descriptor_profile(channels, descriptor.lmax).degree_channels[1]
+    generator = torch.Generator().manual_seed(7)
+
+    def full_rank(knot_data: torch.Tensor) -> torch.Tensor:
+        knot_data = knot_data.clone()
+        knot_data[:, first:channels] = torch.randn(
+            knot_data.shape[0],
+            channels - first,
+            dtype=knot_data.dtype,
+            generator=generator,
+        )
+        return knot_data
+
+    knots = tuple(full_rank(data) for data in _radial_knots(descriptor, 0.002))
+    basis_table, basis_matrix = build_low_rank_table(descriptor, knots=knots)
+    assert basis_table.numel() == 0
+    assert basis_matrix.numel() == 0
+
+
+def test_coarse_knots_keep_the_direct_table() -> None:
+    """Knots too few to determine the basis leave the direct table in place.
+
+    A stride so coarse that the knots give no more samples than the basis has
+    functions would let any channels pass the reproduction check, so no basis
+    is adopted, while the same descriptor adopts one at the default stride.
+    """
+    descriptor = _build_descriptor(64)
+    _restrict_scalar_rank(descriptor, rank=8)
+    assert build_low_rank_table(descriptor)[1].shape[0] == 32
+    basis_table, basis_matrix = build_low_rank_table(
+        descriptor, stride=descriptor.rcut / 8
+    )
+    assert basis_table.numel() == 0
+    assert basis_matrix.numel() == 0
+
+
+def test_unfolded_mode_counts_keep_the_direct_table() -> None:
+    """A descriptor with more modes than the folded scans carry keeps none.
+
+    Its single-type scans read the ordered pair of every edge, so they never
+    read a low-rank table, which compression therefore does not build even
+    for scalar-only channels of low rank.
+    """
+    descriptor = _build_descriptor(64, radial_modes=8)
+    _restrict_scalar_rank(descriptor, rank=8)
+    basis_table, basis_matrix = build_low_rank_table(descriptor)
+    assert basis_table.numel() == 0
+    assert basis_matrix.numel() == 0
diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py
index d1586c5856..a23f8713a6 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py
@@ -2,7 +2,18 @@
 """Numerical contract of the compressed DPA4C CUDA mega kernel."""
 
 import dataclasses
+import functools
+import json
+import math
+import os
+import pathlib
+import re
+import statistics
+import subprocess
+import sys
+import types
 from collections.abc import (
+    Callable,
     Sequence,
 )
 from typing import (
@@ -25,6 +36,7 @@
     DescrptDPA4C,
 )
 from deepmd.pt_expt.kernels.dpa4c.graph_compress import (
+    _portable_inputs,
     _reference_descriptor,
     _reference_forward,
     _table_lookup,
@@ -38,6 +50,10 @@
     op_available,
 )
 
+from .test_dpa4c_cpu import (
+    _restrict_scalar_rank,
+)
+
 _GPU = pytest.mark.skipif(
     not torch.cuda.is_available() or not op_available(),
     reason="CUDA and the compiled DPA4C operator are required",
@@ -49,34 +65,68 @@
 #: against, so a kernel that gathered the wrong radius would misplace it.
 _TYPE_MAP = ["O", "H"]
 
+#: Element symbols of the systems of up to nine types, whose first two are
+#: those of `_TYPE_MAP`. The bridging window `_WINDOW` of every pair of the
+#: first eight stays inside the 3 Å cutoff of the descriptors; that of argon,
+#: the ninth, does not.
+_WIDE_TYPE_MAP = ["O", "H", "C", "N", "B", "F", "Ne", "He", "Ar"]
+
 #: A bridging window wide enough that the random test cluster populates its
 #: frozen zone, its transition zone and the plain zone beyond it. The bounds
 #: are fractions of each pair's own length scale, not distances in Å.
 _WINDOW = {"inner_clamp_f_inner": 0.9, "inner_clamp_f_outer": 1.6}
 
+#: Factor on the mode-mixing columns of the output layer of the pair encoder of
+#: every test descriptor (see `_strengthen_mode_mixing`).
+_MIXING_GAIN = 2000.0
+
+
+def _strengthen_mode_mixing(descriptor: DescrptDPA4C) -> None:
+    """Raise the mode mixing of a fresh descriptor to the scale of its FiLM terms.
+
+    The output layer of a fresh pair encoder gives mixing coefficients of
+    order 1e-4, against FiLM scales of order one, so the mixing of a wrong
+    slot, mode or source type stays below the tolerance of several
+    comparisons, those inside a bridging window among them. Scaling the mixing
+    columns ``[2C, 2C + C R)`` of that layer by ``_MIXING_GAIN`` gives the
+    coefficients of the profiles of 32 channels and more a mean magnitude of
+    0.06 to 0.31 and none above 0.89; those of the profiles of 8 and 16
+    channels reach a mean of 0.47 to 0.61, and in part the saturation of their
+    ``tanh`` bound, which keeps them below one.
+    """
+    film = descriptor.pair_film
+    mixing = slice(2 * film.channels, (2 + film.radial_modes) * film.channels)
+    with torch.no_grad():
+        film.network.layers[-1].w[:, mixing] *= _MIXING_GAIN
+
 
 def _build_descriptor(
     channels: int,
     lmax: int = 2,
     radial_modes: int = 0,
+    rcut: float = 3.0,
+    type_map: Sequence[str] = tuple(_TYPE_MAP),
     **window: float,
 ) -> DescrptDPA4C:
-    return (
+    """Return an fp32 descriptor on the GPU whose mode mixing is strengthened."""
+    descriptor = (
         DescrptDPA4C(
-            rcut=3.0,
-            ntypes=2,
+            rcut=rcut,
+            ntypes=len(type_map),
             channels=channels,
             lmax=lmax,
             n_radial=8,
             radial_modes=radial_modes,
             precision="float32",
             seed=17,
-            type_map=_TYPE_MAP,
+            type_map=list(type_map),
             **window,
         )
         .cuda()
         .eval()
     )
+    _strengthen_mode_mixing(descriptor)
+    return descriptor
 
 
 def _reduced_distance(
@@ -95,18 +145,32 @@ def _build_graph(
     descriptor: DescrptDPA4C,
     canonical: bool,
     node_count: int = 24,
+    single_type: bool = False,
+    side: float = 5.0,
+    frames: int = 1,
+    types: Sequence[int] = (0, 1),
 ):
+    """Return a graph of ``frames`` random clusters of ``node_count`` atoms.
+
+    The atoms of each frame are uniform in a cube of ``side`` Å, and atom ``i``
+    carries type ``types[i % len(types)]``.
+    """
     generator = torch.Generator(device="cuda").manual_seed(23)
     coordinate = torch.rand(
-        1,
+        frames,
         node_count,
         3,
         dtype=torch.float32,
         device="cuda",
         generator=generator,
     )
-    coordinate = coordinate * 5.0
-    atype = torch.arange(node_count, device="cuda").reshape(1, -1) % 2
+    coordinate = coordinate * side
+    pattern = torch.as_tensor(list(types), device="cuda")
+    atype = pattern[torch.arange(node_count, device="cuda") % len(types)].repeat(
+        frames, 1
+    )
+    if single_type:
+        atype = torch.ones_like(atype)
     coord_ext, atype_ext, mapping, nlist = extend_input_and_build_neighbor_list(
         coordinate,
         atype,
@@ -143,6 +207,8 @@ def _arguments(
         graph.destination_row_ptr,
         atype,
         artifacts["data"],
+        artifacts["basis_table"],
+        artifacts["basis_matrix"],
         artifacts["pair_film"],
         artifacts["pair_mixing"],
         artifacts["type_embedding"],
@@ -178,13 +244,19 @@ def _descriptor_backward(
 
 
 def _spin_free(arguments: tuple) -> tuple:
-    """Drop the native spin block for the spin-free CPU reference."""
-    return (*arguments[:15], *arguments[18:])
+    """Return the positional arguments of the reference of a spin-free descriptor.
+
+    The reference (`_reference_descriptor`) reads the radial table itself, so
+    the low-rank table and its expansion are dropped, and takes the native
+    spin block by keyword, so the empty spin block is dropped as well.
+    """
+    portable = _portable_inputs(arguments)
+    return (*portable[:15], *portable[18:])
 
 
 def _with_spin(arguments: tuple, spin: torch.Tensor) -> tuple:
     """Place per-node magnetic moments in the operator's spin slot."""
-    return (*arguments[:15], spin, *arguments[16:])
+    return (*arguments[:17], spin, *arguments[18:])
 
 
 def _assert_dispatched(actual: torch.Tensor, portable: torch.Tensor) -> None:
@@ -202,6 +274,22 @@ def _assert_dispatched(actual: torch.Tensor, portable: torch.Tensor) -> None:
     )
 
 
+def _assert_close_at_scale(
+    actual: torch.Tensor, expected: torch.Tensor, atol: float, rtol: float
+) -> None:
+    """Assert that ``actual`` matches ``expected`` at tolerances stated for unit scale.
+
+    The rounding of an fp32 sum follows the magnitude of its terms, not that of
+    its result, and the terms of a component that cancels to a small value can
+    be as large as the largest components of the output. The absolute
+    tolerance ``atol`` is therefore stated for an output whose largest
+    magnitude is one, and scales with the largest magnitude of ``expected``
+    beyond it.
+    """
+    scale = max(expected.abs().max().item(), 1.0)
+    torch.testing.assert_close(actual, expected, atol=atol * scale, rtol=rtol)
+
+
 @_GPU
 @pytest.mark.parametrize("channels", [8, 16, 32, 64, 128])
 @pytest.mark.parametrize("canonical", [False, True])
@@ -420,6 +508,48 @@ def test_descriptor_compression_routing_and_serialization(
     torch.testing.assert_close(restored_output, actual)
 
 
+@_GPU
+def test_a_snapshot_without_the_low_rank_table_reads_the_direct_table(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A compressed snapshot without the low-rank artifacts reads the radial table.
+
+    A payload that omits ``basis_table`` and ``basis_matrix`` restores both
+    as empty on both backends, which selects the direct table: on a
+    single-type canonical graph, where the scans would read a low-rank table,
+    the restored descriptor evaluates bit for bit as the original with its
+    low-rank artifacts cleared, and differs from the original that reads
+    them. It writes the empty artifacts back.
+    """
+    import copy
+
+    from deepmd.dpmodel.descriptor.dpa4c import DescrptDPA4C as DPDescrptDPA4C
+
+    descriptor = _build_descriptor(64)
+    _restrict_scalar_rank(descriptor, rank=8)
+    descriptor.enable_compression(min_nbor_dist=0.5)
+    data = descriptor.serialize()
+    variables = data["compress"]["@variables"]
+    assert variables["basis_table"].size > 0
+    del variables["basis_table"], variables["basis_matrix"]
+    restored_dpmodel = DPDescrptDPA4C.deserialize(copy.deepcopy(data))
+    assert restored_dpmodel.compress_basis_table.size == 0
+    assert restored_dpmodel.compress_basis_matrix.size == 0
+    restored = DescrptDPA4C.deserialize(copy.deepcopy(data)).cuda().eval()
+
+    low_rank = _check_canonical_against_reference(descriptor, False, monkeypatch)
+    empty = descriptor.compress_basis_table.new_empty(0)
+    descriptor._store_compression_buffer("basis_table", empty)
+    descriptor._store_compression_buffer("basis_matrix", empty.clone())
+    direct = _check_canonical_against_reference(descriptor, False, monkeypatch)
+    actual = _check_canonical_against_reference(restored, False, monkeypatch)
+    assert torch.equal(actual[2], direct[2])
+    assert not torch.equal(low_rank[2], direct[2])
+    written = restored.serialize()["compress"]["@variables"]
+    assert written["basis_table"].size == 0
+    assert written["basis_matrix"].size == 0
+
+
 @_GPU
 def test_compression_uses_immutable_type_snapshot(
     monkeypatch: pytest.MonkeyPatch,
@@ -474,7 +604,8 @@ def _build_charge_descriptor(
 
     The condition output projection is zero initialized, so it is replaced by
     deterministic weights: a parity test against an inert condition would
-    pass for the wrong reason.
+    pass for the wrong reason. The mode mixing is strengthened as in
+    `_build_descriptor`.
     """
     descriptor = (
         DescrptDPA4C(
@@ -507,6 +638,7 @@ def _build_charge_descriptor(
             )
             * 0.5
         )
+    _strengthen_mode_mixing(descriptor)
     return descriptor
 
 
@@ -627,7 +759,8 @@ def test_supported_surface_parity(
     Each scalar width owns a distinct lane mapping and each angular degree a
     distinct instantiation, so the cross product is the contract the operator
     advertises rather than a sample of it. A bridging window selects the
-    bridged specialization of every profile.
+    bridged specialization of every profile. The tolerances follow the
+    magnitude of each output (`_assert_close_at_scale`).
     """
     descriptor = _build_descriptor(
         channels, lmax, radial_modes, **(_WINDOW if windowed else {})
@@ -639,7 +772,121 @@ def test_supported_surface_parity(
     edge_vec = graph.edge_vec.detach().clone().requires_grad_(True)
     output, _state = torch.ops.deepmd.dpa4c_graph_compress(edge_vec, *arguments)
     reference, _ = descriptor.call_graph(graph, atype)
-    torch.testing.assert_close(output, reference, atol=3e-5, rtol=3e-5)
+    _assert_close_at_scale(output, reference, atol=3e-5, rtol=3e-5)
+
+    cotangent = torch.linspace(
+        -0.7,
+        1.3,
+        output.numel(),
+        dtype=output.dtype,
+        device=output.device,
+    ).reshape_as(output)
+    (gradient,) = torch.autograd.grad((output * cotangent).sum(), edge_vec)
+    reference_edge = graph.edge_vec.detach().clone().requires_grad_(True)
+    reference_value = _reference_descriptor(reference_edge, *_spin_free(arguments))
+    (reference_gradient,) = torch.autograd.grad(
+        (reference_value * cotangent).sum(),
+        reference_edge,
+    )
+    # Inside the window the switch is steep, and the float32 reference then
+    # carries a rounding error of its own that exceeds the plain tolerance.
+    atol, rtol = (1e-4, 5e-4) if windowed else (8e-6, 1e-4)
+    _assert_close_at_scale(gradient, reference_gradient, atol=atol, rtol=rtol)
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "type_count", "windowed", "low_rank"),
+    [
+        *(
+            (channels, lmax, radial_modes, 2, False, False)
+            for channels in (32, 64, 128)
+            for lmax in (2, 3)
+            for radial_modes in (2, 4)
+        ),
+        *(
+            (channels, 3, 4, type_count, windowed, False)
+            for channels in (64, 128)
+            for type_count in (3, 4, 5, 8, 9)
+            for windowed in (False, True)
+            if type_count < 9 or not windowed
+        ),
+        (64, 3, 4, 2, True, False),
+        (128, 3, 4, 2, True, False),
+        (64, 4, 4, 3, False, False),
+        (64, 4, 4, 6, True, False),
+        *(
+            (channels, lmax, radial_modes, type_count, windowed, True)
+            for channels, lmax, radial_modes in (
+                (64, 2, 0),
+                (64, 3, 0),
+                (64, 3, 4),
+                (128, 3, 0),
+                (128, 3, 4),
+            )
+            for type_count in (2, 4, 5, 8, 9)
+            for windowed in (False, True)
+            if type_count < 9 or not windowed
+        ),
+    ],
+)
+@pytest.mark.parametrize("system", ["unsplit", "coscheduled"])
+def test_slotted_scans_match_the_portable_equations(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    type_count: int,
+    windowed: bool,
+    low_rank: bool,
+    system: str,
+) -> None:
+    """The scans of a system of a few types follow the portable equations.
+
+    A canonical system of two to eight node types slots the mode profiles and
+    the basis functions of a low-rank table of its edges by the types of their
+    sources, in four slots up to four types and in eight from five, in the edge
+    scans of both directions, except the backward of eight slots of a
+    descriptor without radial modes, which scans pairwise; a system of nine
+    types scans every edge pairwise. Scalar-only channels confined to the span
+    of 32 functions make the compression adopt the low-rank table, every
+    function of which then carries a share of the channels, and which the
+    portable equations do not read. A bridging window selects the bridged
+    specialization. The graph carries 64-bit indices, whose rows always take a
+    warp each; a co-scheduled system holds four waves of single 24-row blocks
+    per multiprocessor, twice as many as the profiles that co-schedule need.
+    The atoms of the random cluster keep the density of the 24-atom one, so its
+    rows hold the short bonds of a random cluster as well. The tolerances
+    follow the magnitude of each output (`_assert_close_at_scale`), which the
+    strengthened mode mixing raises well above one.
+    """
+    multiprocessors = torch.cuda.get_device_properties("cuda").multi_processor_count
+    node_count = {
+        "unsplit": 2 * multiprocessors + 17,
+        "coscheduled": 4 * 24 * multiprocessors + 17,
+    }[system]
+    descriptor = _build_descriptor(
+        channels,
+        lmax,
+        radial_modes,
+        type_map=_WIDE_TYPE_MAP[:type_count],
+        **(_WINDOW if windowed else {}),
+    )
+    if low_rank:
+        _restrict_scalar_rank(descriptor, rank=32)
+    graph, atype = _build_graph(
+        descriptor,
+        canonical=True,
+        node_count=node_count,
+        side=5.0 * (node_count / 24) ** (1.0 / 3.0),
+        types=range(type_count),
+    )
+    arguments = _arguments(descriptor, graph, atype)
+    assert (arguments[6].numel() > 0) is low_rank
+
+    edge_vec = graph.edge_vec.detach().clone().requires_grad_(True)
+    output, _state = torch.ops.deepmd.dpa4c_graph_compress(edge_vec, *arguments)
+    reference, _ = descriptor.call_graph(graph, atype)
+    _assert_close_at_scale(output, reference, atol=3e-5, rtol=3e-5)
 
     cotangent = torch.linspace(
         -0.7,
@@ -658,16 +905,562 @@ def test_supported_surface_parity(
     # Inside the window the switch is steep, and the float32 reference then
     # carries a rounding error of its own that exceeds the plain tolerance.
     atol, rtol = (1e-4, 5e-4) if windowed else (8e-6, 1e-4)
-    torch.testing.assert_close(gradient, reference_gradient, atol=atol, rtol=rtol)
+    _assert_close_at_scale(gradient, reference_gradient, atol=atol, rtol=rtol)
+
+
+@_GPU
+@pytest.mark.parametrize("type_count", [3, 6])
+@pytest.mark.parametrize("windowed", [False, True])
+@pytest.mark.parametrize("perturbation", ["source", "mode"])
+def test_wrong_mode_mixing_fails_the_comparisons(
+    type_count: int,
+    windowed: bool,
+    perturbation: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A scan that applied the wrong mode mixing fails every comparison.
+
+    A slotted scan that misplaced a slot or a mode would mix the mode profiles
+    of an edge with the coefficients of another source type or of another mode.
+    The operator given the mixing table with the source types or with the modes
+    of every pair rotated stands for such a scan, on a system of three types
+    that the scans of four slots serve and on one of six that those of eight
+    slots serve. Its descriptor and its coordinate gradient fail the
+    comparisons of `test_slotted_scans_match_the_portable_equations`, at the
+    wider tolerances of a bridging window, and its forces that of
+    `_check_canonical_against_reference` without a pair potential, which that
+    check also makes for a bridged descriptor: with the strengthened mixing of
+    the test descriptors (`_strengthen_mode_mixing`), such an error exceeds the
+    tolerance of every comparison.
+    """
+    descriptor = _build_descriptor(
+        64,
+        3,
+        4,
+        type_map=_WIDE_TYPE_MAP[:type_count],
+        **(_WINDOW if windowed else {}),
+    )
+    graph, atype = _build_graph(descriptor, canonical=True, types=range(type_count))
+    arguments = _arguments(descriptor, graph, atype)
+    mixing = arguments[9]
+    pairs = mixing.reshape(type_count + 1, type_count + 1, *mixing.shape[1:])
+    wrong = pairs.clone()
+    if perturbation == "source":
+        wrong[:, :type_count] = pairs[:, :type_count].roll(1, dims=1)
+    else:
+        wrong = wrong.roll(1, dims=-1)
+    wrong = wrong.reshape(mixing.shape).contiguous()
+    perturbed = (*arguments[:9], wrong, *arguments[10:])
+
+    edge_vec = graph.edge_vec.detach().clone().requires_grad_(True)
+    output, _ = torch.ops.deepmd.dpa4c_graph_compress(edge_vec, *perturbed)
+    reference, _ = descriptor.call_graph(graph, atype)
+    with pytest.raises(AssertionError):
+        _assert_close_at_scale(output, reference, atol=3e-5, rtol=3e-5)
+    cotangent = torch.linspace(
+        -0.7,
+        1.3,
+        output.numel(),
+        dtype=output.dtype,
+        device=output.device,
+    ).reshape_as(output)
+    (gradient,) = torch.autograd.grad((output * cotangent).sum(), edge_vec)
+    reference_edge = graph.edge_vec.detach().clone().requires_grad_(True)
+    reference_value = _reference_descriptor(reference_edge, *_spin_free(arguments))
+    (reference_gradient,) = torch.autograd.grad(
+        (reference_value * cotangent).sum(),
+        reference_edge,
+    )
+    with pytest.raises(AssertionError):
+        _assert_close_at_scale(gradient, reference_gradient, atol=1e-4, rtol=5e-4)
+
+    artifacts = build_compression_artifacts(descriptor)
+    check = functools.partial(
+        _check_canonical_against_reference,
+        descriptor,
+        False,
+        monkeypatch,
+        single_type=False,
+        types=range(type_count),
+    )
+    descriptor._set_compression(artifacts)
+    forces = check()[2]
+    descriptor._set_compression({**artifacts, "pair_mixing": wrong})
+    wrong_forces = check()[2]
+    scale = max(forces.abs().max().item(), 1.0)
+    with pytest.raises(AssertionError):
+        torch.testing.assert_close(wrong_forces, forces, atol=2e-5 * scale, rtol=2e-5)
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "low_rank"),
+    [
+        (64, 2, 0, True),
+        (64, 3, 4, False),
+        (64, 3, 4, True),
+        (128, 3, 4, False),
+        (128, 3, 4, True),
+    ],
+)
+@pytest.mark.parametrize("system", ["split", "unsplit", "coscheduled"])
+@pytest.mark.parametrize(
+    ("types", "bridged"),
+    [((0, 1), False), ((0, 2, 3, 5, 6, 7), False), ((0, 2, 3, 5, 6, 7), True)],
+)
+def test_fused_slotted_scans_match_the_reference(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    low_rank: bool,
+    system: str,
+    types: tuple[int, ...],
+    bridged: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The fused operator scans systems of a few types slotted, as LAMMPS does.
+
+    The compact canonical graph carries 32-bit indices, whose rows a small
+    system shares among several warps: the split system of 40 atoms shares
+    them, and the larger sizes follow the multiprocessor count so that the rows
+    take a warp each: an unsplit system holds more rows than half a wave of the
+    warps of independent rows, of which a multiprocessor keeps at most 32, and
+    a co-scheduled system holds four waves of single 24-row blocks per
+    multiprocessor. The descriptor has eight types: a system of two takes the
+    scans of four slots, and one of six, which are not consecutive, those of
+    eight, bridged as well. The CPU reference reads the radial table, so the
+    low-rank table of a 32-function basis, which the slotted scans read as the
+    single-type scans do, sums other terms: its forces differ from those of the
+    same descriptor without the basis in the last bits, which shows that the
+    slotted low-rank scans ran.
+    """
+    multiprocessors = torch.cuda.get_device_properties("cuda").multi_processor_count
+    node_count = {
+        "split": 40,
+        "unsplit": 16 * multiprocessors + 17,
+        "coscheduled": 4 * 24 * multiprocessors + 17,
+    }[system]
+    descriptor = _build_descriptor(
+        channels,
+        lmax,
+        radial_modes,
+        type_map=_WIDE_TYPE_MAP[:8],
+        **(_WINDOW if bridged else {}),
+    )
+    if low_rank:
+        _restrict_scalar_rank(descriptor, rank=32)
+    artifacts = build_compression_artifacts(descriptor)
+    assert (artifacts["basis_table"].numel() > 0) is low_rank
+    descriptor._set_compression(artifacts)
+    check = functools.partial(
+        _check_canonical_against_reference,
+        descriptor,
+        bridged,
+        monkeypatch,
+        node_count=node_count,
+        side=5.0 * (node_count / 40) ** (1.0 / 3.0),
+        single_type=False,
+        types=types,
+    )
+    slotted = check()
+    if low_rank:
+        empty = artifacts["basis_table"].new_empty(0)
+        descriptor._set_compression(
+            {**artifacts, "basis_table": empty, "basis_matrix": empty}
+        )
+        direct = check()
+        assert not torch.equal(slotted[2], direct[2])
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "low_rank"),
+    [
+        (64, 2, 0, False),
+        (64, 2, 0, True),
+        (64, 3, 4, False),
+        (64, 3, 4, True),
+        (128, 3, 4, True),
+    ],
+)
+@pytest.mark.parametrize("system", ["split", "unsplit", "coscheduled"])
+@pytest.mark.parametrize("slot_types", [2, 6])
+def test_systems_of_a_few_types_take_the_slotted_scans(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    low_rank: bool,
+    system: str,
+    slot_types: int,
+) -> None:
+    """A system of two to eight node types takes the slotted scans.
+
+    Both systems hold the same atoms of two types and seven isolated atoms,
+    which interact with none: in one the isolated atoms carry the two types of
+    the others, or four other types besides them, a system of two or six types,
+    and in the other the seven remaining types of the descriptor, which make a
+    system of nine types that the slotted scans do not serve. The node counts
+    agree, so every launch choice that follows from them does as well: the rows
+    of the split system of 47 atoms are shared among the warps of a block,
+    those of the larger systems take a warp each, and those of the largest the
+    co-scheduled blocks of the profiles that co-schedule. The slotted scans
+    apply the pair coefficients of each slot to reduced sums, so their forces
+    differ from the pairwise ones in the last bits. A descriptor without radial
+    modes and without a low-rank table has nothing to slot and scans pairwise
+    either way: its forces agree bit for bit.
+    """
+    from deepmd.pt_expt.kernels.dpa4c.canonical import (
+        dpa4c_canonical_compress_energy_force,
+    )
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_neighbor_graph,
+    )
+
+    multiprocessors = torch.cuda.get_device_properties("cuda").multi_processor_count
+    node_count = {
+        "split": 40,
+        "unsplit": 16 * multiprocessors + 17,
+        "coscheduled": 4 * 24 * multiprocessors + 17,
+    }[system]
+    descriptor = _build_descriptor(
+        channels, lmax, radial_modes, type_map=_WIDE_TYPE_MAP
+    )
+    if low_rank:
+        _restrict_scalar_rank(descriptor, rank=32)
+    artifacts = build_compression_artifacts(descriptor)
+    assert (artifacts["basis_table"].numel() > 0) is low_rank
+    descriptor._set_compression(artifacts)
+    neighbor_graph, atype = _build_graph(
+        descriptor,
+        canonical=True,
+        node_count=node_count,
+        side=5.0 * (node_count / 40) ** (1.0 / 3.0),
+    )
+    graph = canonical_graph_from_neighbor_graph(
+        dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node)
+    )
+    # The isolated atoms close both row-pointer arrays with empty rows.
+    isolated = 7
+    extended = dataclasses.replace(
+        graph,
+        n_node=graph.n_node + isolated,
+        n_local=graph.n_local + isolated,
+        destination_row_ptr=torch.cat(
+            [graph.destination_row_ptr] + [graph.destination_row_ptr[-1:]] * isolated
+        ),
+        source_row_ptr=torch.cat(
+            [graph.source_row_ptr] + [graph.source_row_ptr[-1:]] * isolated
+        ),
+    )
+    isolated_types = {2: [0, 1, 0, 1, 0, 1, 0], 6: [2, 3, 4, 5, 0, 1, 0]}[slot_types]
+    few_types = torch.cat([atype, atype.new_tensor(isolated_types)])
+    all_types = torch.cat([atype, atype.new_tensor([2, 3, 4, 5, 6, 7, 8])])
+    fitting = _energy_fitting(descriptor)
+    nodes = atype.shape[0]
+    slotted = dpa4c_canonical_compress_energy_force(
+        descriptor, fitting, extended, few_types, fitting.bias_atom_e[:, 0]
+    )
+    pairwise = dpa4c_canonical_compress_energy_force(
+        descriptor, fitting, extended, all_types, fitting.bias_atom_e[:, 0]
+    )
+    scale = max(pairwise[2].abs().max().item(), 1.0)
+    torch.testing.assert_close(
+        slotted[2][:nodes], pairwise[2][:nodes], atol=2e-5 * scale, rtol=2e-5
+    )
+    assert not slotted[2][nodes:].any()
+    assert not pairwise[2][nodes:].any()
+    slots = radial_modes > 0 or low_rank
+    assert torch.equal(slotted[2][:nodes], pairwise[2][:nodes]) is not slots
+
+
+#: An edge scan of 32-bit indices in a kernel trace: its direction, its rows
+#: per block, its warps per row and its scan follow the index type among its
+#: template arguments, the scan as its value in `deepmd_dpa4c::Scan`.
+_EDGE_SCAN = re.compile(
+    r"(forward|edge_backward)_kernel<[^>]*unsigned int, (\d+), (\d+), "
+    r"\(deepmd_dpa4c::Scan\)(\d)>"
+)
+
+#: Values of `deepmd_dpa4c::Scan` of the single-type scans, which read the
+#: radial or the low-rank table, and of the slotted scans of four and of eight
+#: slots.
+_SINGLE_TYPE_SCANS = {1, 2}
+_SLOTTED_SCANS = {3, 4, 5, 6}
+
+
+def _edge_scans(trace: Any) -> list[tuple[str, int, int, int, float]]:
+    """Return the edge scans of 32-bit indices of a CUDA kernel trace.
+
+    Every entry holds the direction, the rows per block, the warps per row and
+    the scan of one kernel (see `_EDGE_SCAN`), and its duration in µs.
+    """
+    scans = []
+    for event in trace.events():
+        match = _EDGE_SCAN.search(event.name)
+        if match is not None:
+            direction, rows, split, scan = match.groups()
+            duration = event.time_range.elapsed_us()
+            scans.append((direction, int(rows), int(split), int(scan), duration))
+    return scans
+
+
+def _run_in_process(call: str) -> Any:
+    """Return the result of ``call`` in a fresh Python process.
+
+    A check that reads from a CUDA kernel trace which kernels ran evaluates in
+    a process of its own: late in a long process, CUPTI can stop delivering the
+    kernel records, and the trace then shows the launch calls without the
+    kernels. The process imports this module and evaluates ``call``, an
+    expression of its names whose result passes through JSON.
+    """
+    # The module imports from its package, so the process finds it by its
+    # name from the directory above its top package.
+    root = pathlib.Path(__file__).parents[__name__.count(".")]
+    code = (
+        "import importlib, json, sys\n"
+        f"sys.path[:0] = {[str(root), *sys.path]!r}\n"
+        f"module = importlib.import_module({__name__!r})\n"
+        f"print('result:', json.dumps(module.{call}))\n"
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", code], capture_output=True, text=True, check=False
+    )
+    assert result.returncode == 0, result.stdout + "\n" + result.stderr
+    (report,) = [
+        line for line in result.stdout.splitlines() if line.startswith("result:")
+    ]
+    return json.loads(report.removeprefix("result:"))
+
+
+#: Atom counts of the systems of `_row_split_sweep_in_process`: doubling from
+#: one small enough for every edge scan to share its rows among eight warps on
+#: any device.
+_SWEEP_NODE_COUNTS = tuple(8 << k for k in range(13))
+
+
+def _row_split_sweep_in_process(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    rank: int,
+    spin: bool,
+    types: tuple[int, ...],
+) -> list[list]:
+    """Compare the fused operator with its CPU reference on systems of doubling size.
+
+    The systems of `_SWEEP_NODE_COUNTS`, at the density of the 40-atom system
+    and with atoms that cycle through ``types`` of a descriptor of nine types,
+    run until every edge scan takes one warp per row. A positive ``rank``
+    confines the scalar-only channels to the span of that many functions, so
+    that the compression adopts the low-rank table. Returns, per system, its
+    atom count and the sorted direction, scan and warps per row of every edge
+    scan that ran (see `_edge_scans`).
+    """
+    from torch.profiler import (
+        ProfilerActivity,
+        profile,
+    )
+
+    build = _build_spin_descriptor if spin else _build_descriptor
+    descriptor = build(channels, lmax, radial_modes, type_map=_WIDE_TYPE_MAP)
+    if rank:
+        _restrict_scalar_rank(descriptor, rank=rank)
+    artifacts = build_compression_artifacts(descriptor)
+    assert (artifacts["basis_table"].numel() > 0) is bool(rank)
+    descriptor._set_compression(artifacts)
+    sweep = []
+    for node_count in _SWEEP_NODE_COUNTS:
+        with (
+            pytest.MonkeyPatch.context() as patch,
+            profile(activities=[ProfilerActivity.CUDA]) as trace,
+        ):
+            _check_canonical_against_reference(
+                descriptor,
+                False,
+                patch,
+                node_count=node_count,
+                side=5.0 * (node_count / 40) ** (1.0 / 3.0),
+                single_type=False,
+                types=types,
+            )
+        scans = sorted(
+            {
+                (direction, scan, split)
+                for direction, _, split, scan, _ in _edge_scans(trace)
+            }
+        )
+        assert scans, f"the trace of {node_count} atoms records no edge scan"
+        sweep.append([node_count, scans])
+        if all(split == 1 for _, _, split in scans):
+            return sweep
+    raise AssertionError(
+        f"an edge scan still shares its rows at {_SWEEP_NODE_COUNTS[-1]} atoms"
+    )
+
+
+_FEW = (0, 1)
+_SIX = (0, 2, 3, 5, 6, 8)
+_NINE = tuple(range(9))
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "rank", "spin", "types"),
+    [
+        (64, 2, 0, 32, False, _FEW),
+        (64, 3, 4, 32, False, _FEW),
+        (128, 3, 4, 32, False, _FEW),
+        (64, 3, 4, 0, False, _FEW),
+        (64, 3, 4, 0, True, _FEW),
+        (128, 3, 4, 0, True, _FEW),
+        (64, 2, 0, 32, False, _SIX),
+        (64, 3, 4, 32, False, _SIX),
+        (128, 3, 4, 32, False, _SIX),
+        (64, 3, 4, 0, False, _SIX),
+        (64, 3, 4, 0, True, _SIX),
+        (64, 3, 4, 32, False, _NINE),
+        (64, 3, 4, 0, True, _NINE),
+    ],
+)
+def test_rows_shared_among_warps_match_the_reference(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    rank: int,
+    spin: bool,
+    types: tuple[int, ...],
+) -> None:
+    """Rows shared among two, four or eight warps follow the reference.
+
+    An edge scan of a canonical system shares each row among the most of 1, 2,
+    4 and 8 warps whose warps fit in one wave of its resident warps, so once a
+    system is too large for eight warps per row, every doubling of its atoms
+    halves the warps per row. The systems of doubling size of
+    `_row_split_sweep_in_process`, in a process of their own
+    (`_run_in_process`), whose kernel traces show the warps per row of every
+    edge scan, therefore take, in every scan of both directions, each split
+    from that of the smallest system down to one warp per row, and every system
+    follows the reference. The shared rows merge the slot sums of the warps of
+    a row in the slotted scans of four slots of a system of two types and in
+    those of eight slots of a system of six, except for the edge backward of a
+    descriptor without radial modes, which scans six types pairwise, as every
+    profile scans the systems of nine types from sixteen atoms on; the
+    smallest system of that sweep holds eight of the nine types.
+    """
+    sweep = _run_in_process(
+        f"_row_split_sweep_in_process({channels}, {lmax}, {radial_modes}, "
+        f"{rank}, {spin}, {types!r})"
+    )
+    splits: dict[tuple[str, int], list[int]] = {}
+    for _, scans in sweep:
+        for direction, scan, split in scans:
+            splits.setdefault((direction, scan), []).append(split)
+    for kernel, taken in splits.items():
+        assert len(taken) == len(sweep), kernel
+        assert taken[0] > 1, kernel
+        assert taken == sorted(taken, reverse=True), kernel
+        assert set(taken) == {taken[0] >> k for k in range(taken[0].bit_length())}, (
+            kernel
+        )
+
+
+def _padded_storage(graph: Any, slots: int) -> Any:
+    """Return ``graph`` with ``slots`` guard slots past its edge storage.
+
+    The guard slots hold zero bond vectors and the identity entries of the
+    source order, and belong to no row of either CSR, as the guard slots of a
+    storage that a candidate count sizes do.
+    """
+    storage = graph.source.shape[0]
+    order = torch.arange(
+        storage, storage + slots, dtype=torch.int64, device=graph.source.device
+    )
+    return dataclasses.replace(
+        graph,
+        source=torch.cat([graph.source, graph.source.new_zeros(slots)]),
+        edge_vec=torch.cat([graph.edge_vec, graph.edge_vec.new_zeros(slots, 3)]),
+        source_order=torch.cat(
+            [graph.source_order, order.to(graph.source_order.dtype)]
+        ),
+    )
+
+
+@_GPU
+@pytest.mark.parametrize("spin", [False, True])
+def test_bond_prefetch_leaves_every_output_unchanged(spin: bool) -> None:
+    """The bond prefetch leaves every output of the fused operator unchanged.
+
+    The edge scans of a canonical system request the bond vectors of every row,
+    and the slotted scans its source indices, into L2 when the bond vectors of
+    the edge storage exceed the L2 cache. A descriptor of 128 channels, degree
+    three, four radial modes and nine types, with the low-rank table when it is
+    spin-free, evaluates systems of the six types of `_SIX`: one of 40 atoms,
+    whose rows the warps of a block share, one of more rows than half a wave of
+    the warps of independent rows, which take a warp each, and one of four
+    waves of single 24-row blocks per multiprocessor, which takes the
+    co-scheduled blocks. Each system evaluates on its own edge storage, whose
+    bond vectors fit in L2, and on the same storage with guard slots that carry
+    it past the L2 size (`_padded_storage`), and every output of the two agrees
+    bit for bit.
+    """
+    from deepmd.pt_expt.kernels.dpa4c.canonical import (
+        dpa4c_canonical_compress_energy_force,
+    )
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_neighbor_graph,
+    )
+
+    build = _build_spin_descriptor if spin else _build_descriptor
+    descriptor = build(128, 3, 4, type_map=_WIDE_TYPE_MAP)
+    if not spin:
+        _restrict_scalar_rank(descriptor, rank=32)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    fitting = _energy_fitting(descriptor)
+    properties = torch.cuda.get_device_properties("cuda")
+    multiprocessors = properties.multi_processor_count
+    # Bytes of the fp32 bond vector of one storage slot.
+    slot_bytes = 3 * 4
+    prefetch_slots = properties.L2_cache_size // slot_bytes + 1
+    for node_count in (40, 16 * multiprocessors + 17, 4 * 24 * multiprocessors + 17):
+        neighbor_graph, atype = _build_graph(
+            descriptor,
+            canonical=True,
+            node_count=node_count,
+            side=5.0 * (node_count / 40) ** (1.0 / 3.0),
+            types=_SIX,
+        )
+        graph = canonical_graph_from_neighbor_graph(
+            dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node)
+        )
+        storage = graph.source.shape[0]
+        assert storage < prefetch_slots, node_count
+        moments = None
+        if spin:
+            generator = torch.Generator(device="cuda").manual_seed(13)
+            moments = torch.randn(
+                atype.shape[0],
+                3,
+                dtype=torch.float32,
+                device="cuda",
+                generator=generator,
+            )
+        outputs = [
+            dpa4c_canonical_compress_energy_force(
+                descriptor, fitting, edges, atype, fitting.bias_atom_e[:, 0], moments
+            )
+            for edges in (graph, _padded_storage(graph, prefetch_slots - storage))
+        ]
+        for index, (compact, padded) in enumerate(zip(*outputs, strict=True)):
+            assert torch.equal(compact, padded), (node_count, index)
 
 
 @_GPU
 @pytest.mark.parametrize(
     ("index", "shape", "dtype", "message"),
     [
-        (6, (9, 4, 2), torch.float32, "PairFiLM"),
-        (9, (8, 5, 5), torch.float32, "invalid readout matrix"),
-        (10, (1, 8), torch.int32, "degree triples"),
+        (8, (9, 4, 2), torch.float32, "PairFiLM"),
+        (11, (8, 5, 5), torch.float32, "invalid readout matrix"),
+        (12, (1, 8), torch.int32, "degree triples"),
     ],
 )
 def test_operator_rejects_inconsistent_artifacts(
@@ -685,6 +1478,70 @@ def test_operator_rejects_inconsistent_artifacts(
         torch.ops.deepmd.dpa4c_graph_compress(graph.edge_vec, *arguments)
 
 
+def _misaligned(tensor: torch.Tensor) -> torch.Tensor:
+    """Return a copy of ``tensor`` that starts four bytes past its alignment."""
+    storage = torch.zeros(tensor.numel() + 1, dtype=tensor.dtype, device=tensor.device)
+    copy = storage[1:].view(tensor.shape)
+    copy.copy_(tensor)
+    return copy
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "case", "message"),
+    [
+        (8, "atype", "atype must be one-dimensional"),
+        (8, "stride", "table_stride must be positive and finite"),
+        (8, "reach", "table_max must cover rcut within"),
+        (8, "basis", "8 channels admit no low-rank radial table"),
+        (8, "film", "8-byte aligned"),
+        (64, "basis_width", "invalid low-rank radial table shape"),
+        (64, "expansion_shape", "the low-rank expansion must have shape"),
+        (64, "expansion_alone", "a low-rank expansion requires its table"),
+        (64, "expansion_alignment", "16-byte aligned"),
+    ],
+)
+def test_operator_rejects_invalid_launch_inputs(
+    channels: int, case: str, message: str
+) -> None:
+    """The operator refuses launch inputs that the scans cannot read.
+
+    A two-dimensional atom-type tensor, a stride that is not finite, a table
+    reach beyond the knots of the radial table, a low-rank table for a profile
+    without the expansion or of the wrong width, an expansion of the wrong
+    shape or without its table, and tables off the alignment of the vector
+    loads of the scans each fail with their message. The descriptor of 64
+    channels carries the low-rank table of `_restrict_scalar_rank`.
+    """
+    descriptor = _build_descriptor(channels)
+    if channels == 64:
+        _restrict_scalar_rank(descriptor, rank=32)
+    graph, atype = _build_graph(descriptor, canonical=False)
+    arguments = list(_arguments(descriptor, graph, atype))
+    table, basis_table, basis_matrix = arguments[5:8]
+    stride = arguments[23]
+    if case == "atype":
+        arguments[4] = atype[:, None]
+    elif case == "stride":
+        arguments[23] = math.nan
+    elif case == "reach":
+        arguments[24] = 1.01 * table.shape[0] * stride
+    elif case == "basis":
+        arguments[6] = torch.zeros_like(table)
+    elif case == "film":
+        arguments[8] = _misaligned(arguments[8])
+    elif case == "basis_width":
+        arguments[6] = basis_table[:, 6:].contiguous()
+    elif case == "expansion_shape":
+        arguments[7] = basis_matrix[1:].contiguous()
+    elif case == "expansion_alone":
+        arguments[6] = basis_table[:0]
+    else:
+        arguments[7] = _misaligned(basis_matrix)
+    with pytest.raises(RuntimeError, match=message):
+        torch.ops.deepmd.dpa4c_graph_compress(graph.edge_vec, *arguments)
+
+
 @_GPU
 def test_operator_rejects_unsupported_mode_rank() -> None:
     """An unsupported mode rank would overflow the shared mode cache."""
@@ -701,8 +1558,8 @@ def test_operator_rejects_unsupported_mode_rank() -> None:
         dtype=torch.float32,
         device="cuda",
     )
-    arguments[7] = torch.zeros(
-        arguments[7].shape[0],
+    arguments[9] = torch.zeros(
+        arguments[9].shape[0],
         channels,
         3,
         dtype=torch.float32,
@@ -1104,8 +1961,8 @@ def test_compact_canonical_parity(
         graph.edge_index[0].to(index_dtype),
         graph.destination_row_ptr,
         atype,
-        *arguments[5:19],
-        *arguments[20:],
+        *arguments[5:21],
+        *arguments[22:],
     )
     compact_output, compact_state = torch.ops.deepmd.dpa4c_canonical_compress(
         graph.edge_vec,
@@ -1149,8 +2006,8 @@ def test_compact_inplace_backward_reuses_state(channels: int) -> None:
         graph.edge_index[0].to(torch.uint32),
         graph.destination_row_ptr,
         atype,
-        *arguments[5:19],
-        *arguments[20:],
+        *arguments[5:21],
+        *arguments[22:],
     )
     ensure_canonical_registered()
     output, state = torch.ops.deepmd.dpa4c_canonical_compress(
@@ -1322,17 +2179,59 @@ def test_fused_energy_force_refuses_ineligible_fitting() -> None:
         )
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-@pytest.mark.parametrize("channels", [8, 64])
-def test_compact_canonical_tiling_is_equivalent(
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "low_rank", "system"),
+    [
+        (8, 2, 0, False, "split"),
+        (64, 2, 0, False, "split"),
+        (64, 2, 0, False, "unsplit"),
+        (128, 2, 0, False, "coscheduled"),
+        (64, 3, 4, False, "split"),
+        (64, 3, 4, False, "unsplit"),
+        (64, 2, 0, True, "split"),
+        (64, 2, 0, True, "unsplit"),
+        (128, 3, 4, True, "coscheduled"),
+    ],
+)
+@pytest.mark.parametrize("type_count", [2, 6])
+def test_compact_canonical_tiling_is_equivalent(
     channels: int,
+    lmax: int,
+    radial_modes: int,
+    low_rank: bool,
+    system: str,
+    type_count: int,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """Node tiling must not change energy, force or virial.
-
-    Every node tile owns a contiguous span of the destination-sorted edge
-    axis, so the runs partition the work rather than splitting any reduction.
+    """Node tiling changes no output in any bit.
+
+    Every node tile owns a contiguous span of the destination-sorted edge axis,
+    so the runs partition the work rather than splitting any reduction, and
+    every run scans its rows with the schedule of the whole system, whose row
+    split, co-scheduled blocks and slotted scans a run of other length or of
+    fewer node types would not take. The rows of the graph are sorted by bond
+    length, so the lockstep windows of a co-scheduled row follow from the row
+    alone, whichever rows share its block. The sizes follow the multiprocessor
+    count: an unsplit system holds more rows than half a wave of the warps of
+    independent rows, of which a multiprocessor keeps at most 32, so it scans
+    one warp per row while its shorter runs would share each row; a
+    co-scheduled system holds four waves of single 24-row blocks per
+    multiprocessor, which its runs of less than one wave would not take. The
+    system holds two or six types, so the profiles with radial modes or a
+    low-rank table scan it slotted, in four or eight slots, but for the edge
+    backward of six types of a profile without radial modes, which scans
+    pairwise. The atoms keep the density of the smallest system.
     """
+    multiprocessors = torch.cuda.get_device_properties("cuda").multi_processor_count
+    node_count, tiles = {
+        "split": (97, (7, 32, 96)),
+        "unsplit": (16 * multiprocessors + 17, (100, 4 * multiprocessors)),
+        "coscheduled": (
+            4 * 24 * multiprocessors + 17,
+            (12 * multiprocessors, 24 * multiprocessors),
+        ),
+    }[system]
     from deepmd.pt_expt.fitting.ener_fitting import (
         EnergyFittingNet,
     )
@@ -1343,15 +2242,27 @@ def test_compact_canonical_tiling_is_equivalent(
         canonical_graph_from_neighbor_graph,
     )
 
-    descriptor = _build_descriptor(channels)
-    neighbor_graph, atype = _build_graph(descriptor, canonical=True, node_count=97)
+    descriptor = _build_descriptor(
+        channels, lmax, radial_modes, type_map=_WIDE_TYPE_MAP[:type_count]
+    )
+    if low_rank:
+        _restrict_scalar_rank(descriptor, rank=8)
+    neighbor_graph, atype = _build_graph(
+        descriptor,
+        canonical=True,
+        node_count=node_count,
+        side=5.0 * (node_count / 97) ** (1.0 / 3.0),
+        types=range(type_count),
+    )
     graph = canonical_graph_from_neighbor_graph(
         dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node)
     )
-    descriptor._set_compression(build_compression_artifacts(descriptor))
+    artifacts = build_compression_artifacts(descriptor)
+    assert (artifacts["basis_table"].numel() > 0) is low_rank
+    descriptor._set_compression(artifacts)
     fitting = (
         EnergyFittingNet(
-            ntypes=2,
+            ntypes=type_count,
             dim_descrpt=descriptor.get_dim_out(),
             neuron=[64, 64],
             resnet_dt=False,
@@ -1363,7 +2274,6 @@ def test_compact_canonical_tiling_is_equivalent(
         .cuda()
         .eval()
     )
-    ownership = torch.ones(atype.shape[0], dtype=torch.bool, device="cuda")
 
     def run() -> tuple[torch.Tensor, ...]:
         return dpa4c_canonical_compress_energy_force(
@@ -1371,23 +2281,24 @@ def run() -> tuple[torch.Tensor, ...]:
             fitting,
             graph,
             atype,
-            ownership,
             fitting.bias_atom_e[:, 0],
-            True,
         )
 
     monkeypatch.setenv("DP_NODE_TILE", "0")
     reference = run()
-    for tile in ("7", "32", "96"):
-        monkeypatch.setenv("DP_NODE_TILE", tile)
+    for tile in tiles:
+        monkeypatch.setenv("DP_NODE_TILE", str(tile))
         for actual, expected in zip(run(), reference, strict=True):
-            torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-6)
+            assert torch.equal(actual, expected), f"tile {tile}"
 
 
 @_GPU
 @pytest.mark.parametrize("bridged", [False, True])
+@pytest.mark.parametrize(("frames", "halo"), [(1, 0), (2, 0), (2, 3)])
 def test_fused_canonical_cpu_reference_matches_the_kernel(
     bridged: bool,
+    frames: int,
+    halo: int,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """The CPU implementation of the fused canonical operator reproduces it.
@@ -1395,17 +2306,85 @@ def test_fused_canonical_cpu_reference_matches_the_kernel(
     The operator runs once on CUDA through the energy-force helper and once on
     CPU copies of the same arguments, where it dispatches to the reference
     sequence of descriptor, fitting, fitting backward and descriptor backward,
-    with the pair potential when the model is bridged.
+    with the pair potential when the model is bridged. A batch of two frames
+    takes the multi-frame path of the operator, its frame lookup and its
+    segment sums, with the last ``halo`` nodes of each frame unowned.
     """
-    from deepmd.dpmodel.atomic_model.inner_potential import (
-        InnerPotential,
-    )
-    from deepmd.pt_expt.fitting.ener_fitting import (
-        EnergyFittingNet,
-    )
-    from deepmd.pt_expt.kernels.dpa4c import (
-        canonical,
+    window = {"inner_clamp_f_inner": 0.9, "inner_clamp_f_outer": 1.3} if bridged else {}
+    descriptor = _build_descriptor(8, **window)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    _check_canonical_against_reference(
+        descriptor,
+        bridged,
+        monkeypatch,
+        node_count=24,
+        single_type=False,
+        halo=halo,
+        frames=frames,
     )
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "bridged"),
+    [
+        (32, 2, 0, False),
+        (64, 2, 0, True),
+        (64, 3, 4, False),
+        (128, 3, 4, True),
+        (64, 3, 2, False),
+        (64, 2, 8, False),
+    ],
+)
+def test_single_type_canonical_matches_the_reference(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    bridged: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A single-type system matches the CPU reference of the fused operator.
+
+    Every edge of such a system carries one ordered pair, which the CUDA scans
+    exploit by applying the PairFiLM and mode coefficients to reduced sums and
+    folded cotangents instead of to every edge. Eight modes exceed what the
+    folded scans carry and pin the unfolded path of the same system.
+    """
+    window = {"inner_clamp_f_inner": 0.9, "inner_clamp_f_outer": 1.3} if bridged else {}
+    descriptor = _build_descriptor(channels, lmax, radial_modes, **window)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    _check_canonical_against_reference(descriptor, bridged, monkeypatch)
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "folds"),
+    [(64, 2, 0, True), (128, 3, 4, True), (64, 2, 8, False)],
+)
+@pytest.mark.parametrize("spin", [False, True])
+def test_single_type_systems_take_the_folded_scans(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    folds: bool,
+    spin: bool,
+) -> None:
+    """A single-type system takes the folded scans, which sum other terms.
+
+    Both systems hold the same atoms and one isolated atom, which interacts
+    with none; in one of them it carries the type of the others and in the
+    other the second type, which ends the single-type system, so the scans
+    read the ordered pair of every edge again, or slot the mode profiles of
+    the two types for a profile that carries modes. The node counts agree, so
+    every launch choice that follows from them does as well. The folded scans
+    apply the pair coefficients to reduced sums once per node, so their forces
+    differ from the pairwise and the slotted ones in the last bits. Eight
+    modes exceed what the folded and the slotted scans carry: that profile
+    scans pairwise either way, and its forces agree bit for bit. A
+    spin-conditioned profile has no single-type scan and scans both systems
+    alike, pairwise or slotted, and its forces and magnetic forces agree bit
+    for bit; its atoms are of type 0, which carries the spin branch.
+    """
     from deepmd.pt_expt.kernels.dpa4c.canonical import (
         dpa4c_canonical_compress_energy_force,
     )
@@ -1413,207 +2392,847 @@ def test_fused_canonical_cpu_reference_matches_the_kernel(
         canonical_graph_from_neighbor_graph,
     )
 
-    canonical.ensure_registered()
-    window = {"inner_clamp_f_inner": 0.9, "inner_clamp_f_outer": 1.3} if bridged else {}
-    descriptor = _build_descriptor(8, **window)
-    neighbor_graph, atype = _build_graph(descriptor, canonical=True)
+    build = _build_spin_descriptor if spin else _build_descriptor
+    descriptor = build(channels, lmax, radial_modes)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    neighbor_graph, atype = _build_graph(
+        descriptor, canonical=True, node_count=40, single_type=True
+    )
     graph = canonical_graph_from_neighbor_graph(
         dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node)
     )
-    descriptor._set_compression(build_compression_artifacts(descriptor))
-    fitting = (
-        EnergyFittingNet(
-            ntypes=2,
-            dim_descrpt=descriptor.get_dim_out(),
-            neuron=[32, 32],
-            resnet_dt=False,
-            activation_function="silu",
-            precision="float32",
-            mixed_types=True,
-            seed=29,
+    system_type, isolated_type = (0, 1) if spin else (1, 0)
+    atype = torch.full_like(atype, system_type)
+    moments = None
+    if spin:
+        generator = torch.Generator(device="cuda").manual_seed(13)
+        moments = torch.randn(
+            atype.shape[0] + 1,
+            3,
+            dtype=torch.float32,
+            device="cuda",
+            generator=generator,
         )
-        .cuda()
-        .eval()
+    # The isolated atom closes both row-pointer arrays with an empty row.
+    extended = dataclasses.replace(
+        graph,
+        n_node=graph.n_node + 1,
+        n_local=graph.n_local + 1,
+        destination_row_ptr=torch.cat(
+            [graph.destination_row_ptr, graph.destination_row_ptr[-1:]]
+        ),
+        source_row_ptr=torch.cat([graph.source_row_ptr, graph.source_row_ptr[-1:]]),
     )
-    pair_table = (
-        torch.as_tensor(InnerPotential(["O", "H"]).pair_table, device="cuda")
-        if bridged
-        else None
+    uniform_type = torch.cat([atype, torch.full_like(atype[:1], system_type)])
+    mixed_type = torch.cat([atype, torch.full_like(atype[:1], isolated_type)])
+    fitting = _energy_fitting(descriptor)
+    nodes = atype.shape[0]
+    single = dpa4c_canonical_compress_energy_force(
+        descriptor, fitting, extended, uniform_type, fitting.bias_atom_e[:, 0], moments
     )
-    operator = torch.ops.deepmd.dpa4c_canonical_compress_energy_gradient
-    calls = []
-
-    def record(*args: Any) -> tuple[torch.Tensor, ...]:
-        calls.append(args)
-        return operator(*args)
-
-    monkeypatch.setattr(
-        torch.ops.deepmd, "dpa4c_canonical_compress_energy_gradient", record
+    paired = dpa4c_canonical_compress_energy_force(
+        descriptor, fitting, extended, mixed_type, fitting.bias_atom_e[:, 0], moments
     )
-    dpa4c_canonical_compress_energy_force(
-        descriptor,
-        fitting,
-        graph,
-        atype,
-        torch.ones(atype.shape[0], dtype=torch.bool, device="cuda"),
-        fitting.bias_atom_e[:, 0],
-        False,
-        None,
-        pair_table,
+    scale = max(single[2].abs().max().item(), 1.0)
+    torch.testing.assert_close(
+        paired[2][:nodes], single[2][:nodes], atol=2e-5 * scale, rtol=2e-5
     )
-    (arguments,) = calls
-
-    def to_cpu(value: Any) -> Any:
-        if isinstance(value, torch.Tensor):
-            return value.cpu()
-        if isinstance(value, list):
-            return [to_cpu(item) for item in value]
-        return value
-
-    kernel = operator(*arguments)
-    reference = operator(*[to_cpu(argument) for argument in arguments])
-    for actual, expected in zip(reference, kernel, strict=True):
-        scale = max(expected.abs().max().item(), 1.0) if expected.numel() else 1.0
-        torch.testing.assert_close(actual, expected.cpu(), atol=2e-5 * scale, rtol=2e-5)
+    torch.testing.assert_close(
+        paired[1][:nodes], single[1][:nodes], atol=2e-5, rtol=2e-5
+    )
+    assert not paired[2][nodes:].any()
+    assert not single[2][nodes:].any()
+    folded = folds and not spin
+    assert torch.equal(paired[2][:nodes], single[2][:nodes]) is not folded
+    if spin:
+        assert torch.equal(paired[5][:nodes], single[5][:nodes])
 
 
-def _build_spin_descriptor(
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "bridged"),
+    [(32, 2, 0, False), (64, 2, 0, True), (64, 3, 4, False), (128, 3, 4, True)],
+)
+def test_single_type_spin_canonical_matches_the_reference(
     channels: int,
-    lmax: int = 2,
-    radial_modes: int = 0,
-    device: str = "cuda",
-    **window: float,
-) -> DescrptDPA4C:
-    """Return a spin-conditioned descriptor with a non-unit reference moment.
+    lmax: int,
+    radial_modes: int,
+    bridged: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A single-type spin system matches the CPU reference of the fused operator.
 
-    A reference magnitude other than one makes the conditioning factor visible
-    in the magnetic gradient, so a missing chain factor shows up as a constant
-    ratio rather than cancelling. The branch gate is opened, since a fresh
-    descriptor starts spin-free by design and these tests are about the branch
-    behind the gate; the gate itself is covered in ``test_dpa4c.py``.
+    A spin-conditioned descriptor has no single-type scan: its single-type
+    system takes the slotted scans of a profile with radial modes, and the
+    pairwise scans of any other; energy, force, virial and magnetic force
+    follow the reference.
     """
-    descriptor = (
-        DescrptDPA4C(
-            rcut=3.0,
-            ntypes=2,
-            channels=channels,
-            lmax=lmax,
-            n_radial=8,
-            radial_modes=radial_modes,
-            precision="float32",
-            seed=17,
-            type_map=_TYPE_MAP,
-            use_spin=[True, False],
-            **window,
-        )
-        .to(device)
-        .eval()
-    )
-    descriptor.spin.set_spin_reference(np.array([1.7, 1.0, 1.0]))
-    with torch.no_grad():
-        descriptor.spin.spin_gate.fill_(1.0)
-    return descriptor
+    window = {"inner_clamp_f_inner": 0.9, "inner_clamp_f_outer": 1.3} if bridged else {}
+    descriptor = _build_spin_descriptor(channels, lmax, radial_modes, **window)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    _check_canonical_against_reference(descriptor, bridged, monkeypatch)
 
 
-def _empty_graph(device: str) -> tuple[NeighborGraph, torch.Tensor]:
-    """Return a native graph whose node and edge axes are both empty."""
-    n_node = torch.zeros(1, dtype=torch.int64, device=device)
-    edge_index = torch.empty(2, 0, dtype=torch.int64, device=device)
-    edge_vec = torch.empty(0, 3, dtype=torch.float32, device=device)
-    edge_order = torch.empty(0, dtype=torch.int64, device=device)
-    row_pointer = torch.zeros(1, dtype=torch.int64, device=device)
-    return (
-        NeighborGraph(
-            n_node=n_node,
-            n_local=n_node,
-            edge_index=edge_index,
-            edge_vec=edge_vec,
-            edge_mask=torch.empty(0, dtype=torch.bool, device=device),
-            destination_order=edge_order,
-            destination_row_ptr=row_pointer,
-            source_order=edge_order,
-            source_row_ptr=row_pointer,
-        ),
-        torch.empty(0, dtype=torch.int64, device=device),
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "bridged"),
+    [(64, 2, 0, False), (64, 3, 4, True), (128, 3, 4, False)],
+)
+def test_low_rank_table_matches_the_reference(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    bridged: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The low-rank single-type scans match the CPU reference.
+
+    Scalar-only channels confined to few functions are reproduced by the
+    basis, so the compression adopts the low-rank table, which the single-type
+    scans read in place of the radial table, while the reference evaluates the
+    radial table itself. The low-rank scans sum other terms than the direct
+    ones, so their forces differ from those of the same descriptor without the
+    basis in the last bits, which shows that they ran.
+    """
+    window = {"inner_clamp_f_inner": 0.9, "inner_clamp_f_outer": 1.3} if bridged else {}
+    descriptor = _build_descriptor(channels, lmax, radial_modes, **window)
+    _restrict_scalar_rank(descriptor, rank=8)
+    artifacts = build_compression_artifacts(descriptor)
+    first = descriptor_profile(channels, lmax).degree_channels[1]
+    assert artifacts["basis_matrix"].shape == (32, channels - first)
+    assert artifacts["basis_table"].shape == (
+        artifacts["data"].shape[0],
+        6 * (first + 32 + radial_modes),
     )
-
-
-def test_empty_native_spin_cpu_profile_preserves_spin_contract() -> None:
-    """The CPU reference keeps the native-spin state width at zero nodes."""
-    descriptor = _build_spin_descriptor(8, device="cpu")
-    graph, atype = _empty_graph("cpu")
-    arguments = _with_spin(
-        _arguments(descriptor, graph, atype),
-        graph.edge_vec,
+    descriptor._set_compression(artifacts)
+    low_rank = _check_canonical_against_reference(descriptor, bridged, monkeypatch)
+    empty = artifacts["basis_table"].new_empty(0)
+    descriptor._set_compression(
+        {**artifacts, "basis_table": empty, "basis_matrix": empty}
     )
-
-    output, state = _reference_forward(graph.edge_vec, *arguments)
-
-    profile = descriptor_profile(8, 2, True)
-    assert output.shape == (0, profile.output_width)
-    assert state.shape == (0, profile.state_width)
+    direct = _check_canonical_against_reference(descriptor, bridged, monkeypatch)
+    assert not torch.equal(low_rank[2], direct[2])
 
 
 @_GPU
-def test_empty_native_spin_cuda_backward_preserves_spin_contract() -> None:
-    """The CUDA backward distinguishes present spin from its empty axis."""
-    descriptor = _build_spin_descriptor(8)
-    graph, atype = _empty_graph("cuda")
-    arguments = _with_spin(
-        _arguments(descriptor, graph, atype),
-        graph.edge_vec,
+@pytest.mark.parametrize("single_type", [True, False])
+@pytest.mark.parametrize("placeholders", [False, True])
+def test_canonical_halo_nodes_match_the_reference(
+    single_type: bool,
+    placeholders: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Unowned trailing nodes and source-CSR placeholders match the reference.
+
+    The nodes past ``n_local`` lend their bonds to the owned nodes but carry
+    no energy of their own, and the two-entry placeholders that stand in for
+    the source CSR of a caller that declines the per-atom virial leave that
+    output zero.
+    """
+    descriptor = _build_descriptor(
+        64, 3, 4, inner_clamp_f_inner=0.9, inner_clamp_f_outer=1.3
     )
-    output, state = torch.ops.deepmd.dpa4c_graph_compress(
-        graph.edge_vec,
-        *arguments,
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    kernel = _check_canonical_against_reference(
+        descriptor,
+        True,
+        monkeypatch,
+        single_type=single_type,
+        halo=9,
+        placeholders=placeholders,
     )
-
-    edge_gradient, spin_gradient, edge_spin_gradient = _descriptor_backward(
-        torch.empty_like(output),
-        state,
-        graph.edge_vec,
-        *arguments,
+    assert not kernel[6][-9:].any()
+    if placeholders:
+        assert not kernel[4].any()
+
+
+def _device_filling_scans(
+    single_type: bool, spin: bool
+) -> list[tuple[str, int, int, int, float]]:
+    """Compare the fused operator with its CPU reference on a device-filling system.
+
+    The system holds two blocks of 24 rows per multiprocessor and 17 atoms
+    more, of one type or of two, at a density of 0.08 atoms per cubic Å, which
+    puts about 40 atoms within the 5 Å cutoff of an atom inside the cube. Its
+    descriptor has 128 channels and degree two, and four radial modes when it
+    is spin-free and none when it is spin-conditioned. Returns the edge scans
+    that the comparison launched (see `_edge_scans`).
+    """
+    from torch.profiler import (
+        ProfilerActivity,
+        profile,
     )
 
-    assert edge_gradient.shape == (0, 3)
-    assert spin_gradient.shape == (0, 3)
-    assert edge_spin_gradient.shape == (0, 3)
+    multiprocessors = torch.cuda.get_device_properties("cuda").multi_processor_count
+    node_count = 2 * 24 * multiprocessors + 17
+    if spin:
+        descriptor = _build_spin_descriptor(128, 2, 0, rcut=5.0)
+    else:
+        descriptor = _build_descriptor(128, 2, 4, rcut=5.0)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    with (
+        pytest.MonkeyPatch.context() as patch,
+        profile(activities=[ProfilerActivity.CUDA]) as trace,
+    ):
+        _check_canonical_against_reference(
+            descriptor,
+            False,
+            patch,
+            node_count=node_count,
+            side=(node_count / 0.08) ** (1.0 / 3.0),
+            single_type=single_type,
+        )
+    return _edge_scans(trace)
 
 
 @_GPU
 @pytest.mark.parametrize(
-    ("channels", "lmax", "radial_modes"),
-    [(8, 2, 0), (16, 2, 0), (32, 2, 4), (64, 3, 0), (128, 3, 8), (32, 4, 0)],
+    ("single_type", "spin"), [(False, False), (True, False), (False, True)]
 )
-@pytest.mark.parametrize("windowed", [False, True])
-def test_spin_compressed_matches_portable(
-    monkeypatch: pytest.MonkeyPatch,
-    channels: int,
-    lmax: int,
-    radial_modes: int,
-    windowed: bool,
+def test_coscheduled_scans_match_the_reference(
+    single_type: bool,
+    spin: bool,
 ) -> None:
-    """Descriptor, coordinate gradient and magnetic force all match.
+    """Co-scheduled edge scans of a device-filling graph match the reference.
+
+    A system of at least two blocks of 24 co-scheduled rows per multiprocessor
+    (`_device_filling_scans`) walks the rows of a wide profile in lockstep
+    windows of 4 Å of bond length, and a cutoff beyond one window makes every
+    block realign at a barrier inside its rows. The small systems of the other
+    tests leave every row to warps of its own instead. The kernel trace shows
+    that both scans ran with co-scheduled blocks in the kernels that serve the
+    system: the longest edge scan of each direction is co-scheduled, and is a
+    slotted scan for a system of two types and a single-type scan for one of
+    one type, whose launches also start the kernels of the other, whose blocks
+    retire at once. A spin profile without radial modes has nothing to slot and
+    scans its system of two types pairwise, in the kernels of the single-type
+    scans. The node count gives every multiprocessor two blocks of the widest
+    co-scheduled shape on any device. The comparison and its kernel trace run
+    in a process of their own (`_run_in_process`).
+    """
+    scans = _run_in_process(f"_device_filling_scans({single_type}, {spin})")
+    assert scans
+    # Blocks of more than the four independent rows of a slotted or a spin
+    # scan are co-scheduled ones.
+    slotted = not (single_type or spin)
+    serving = _SLOTTED_SCANS if slotted else _SINGLE_TYPE_SCANS
+    for direction in ("forward", "edge_backward"):
+        _, rows, _, scan, _ = max(
+            (entry for entry in scans if entry[0] == direction),
+            key=lambda entry: entry[4],
+        )
+        assert rows > 4, direction
+        assert scan in serving, direction
 
-    The portable path is the oracle for the spin families. Probing every
-    output column at once keeps the geometric and the spin blocks in the same
-    comparison, because the two are coupled through the shared normalizer and
-    through the cross Gram. The cotangent stays bounded: weighting columns by
-    their index inflates the gradient magnitude at the widest profile until
-    fp32 tabulation noise alone exceeds the tolerance.
+
+@_GPU
+@pytest.mark.parametrize("frames", [1, 2])
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+@pytest.mark.parametrize("spin", [False, True])
+def test_non_finite_energy_marks_the_forces_of_its_frame(
+    frames: int, device: str, spin: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A frame whose energy is not finite reports its forces as not a number.
+
+    A fixed-point force word cannot carry a non-finite value, so the operator
+    marks the forces of every frame whose totals are not finite, and only
+    those, together with their magnetic forces in a spin-conditioned model.
+    The CPU reference of the operator marks the same forces.
     """
-    descriptor = _build_spin_descriptor(
-        channels, lmax, radial_modes, **(_WINDOW if windowed else {})
-    )
-    assert mega_eligible(descriptor)
-    graph, atype = _build_graph(descriptor, canonical=True)
-    generator = torch.Generator(device="cuda").manual_seed(11)
-    spin = torch.randn(
-        atype.shape[0],
-        3,
-        dtype=torch.float32,
-        device="cuda",
-        generator=generator,
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_neighbor_graph,
+    )
+
+    build = _build_spin_descriptor if spin else _build_descriptor
+    descriptor = build(64)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    neighbor_graph, atype = _build_graph(
+        descriptor, canonical=True, node_count=40, frames=frames
+    )
+    graph = canonical_graph_from_neighbor_graph(
+        dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node)
+    )
+    fitting = _energy_fitting(descriptor)
+    # The first frame consists of the type whose energy bias is not a number,
+    # which is the type that carries the spin branch.
+    frame = torch.repeat_interleave(
+        torch.arange(frames, device="cuda"), neighbor_graph.n_node
+    )
+    atype = (frame != 0).to(atype.dtype)
+    atom_bias = torch.tensor([math.nan, 0.0], dtype=torch.float64, device="cuda")
+    moments = None
+    if spin:
+        generator = torch.Generator(device="cuda").manual_seed(13)
+        moments = torch.randn(
+            atype.shape[0], 3, dtype=torch.float32, device="cuda", generator=generator
+        )
+    operator = torch.ops.deepmd.dpa4c_canonical_compress_energy_force
+    arguments = _fused_arguments(
+        monkeypatch, descriptor, fitting, graph, atype, atom_bias, moments
+    )
+    if device == "cpu":
+        arguments = [_to_cpu(argument) for argument in arguments]
+    energy, _, force, _, _, force_mag, _ = operator(*arguments)
+    broken = (frame == 0).to(force.device)
+    assert torch.isnan(energy[0]).all()
+    assert torch.isfinite(energy[1:]).all()
+    assert torch.isnan(force[broken]).all()
+    assert torch.isfinite(force[~broken]).all()
+    if spin:
+        assert torch.isnan(force_mag[broken]).all()
+        assert torch.isfinite(force_mag[~broken]).all()
+
+
+@_GPU
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+def test_an_incomplete_source_csr_is_refused(
+    device: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A source CSR is complete or two-entry placeholders, and nothing else.
+
+    The operator and its CPU reference refuse a source CSR of any other
+    extent rather than reading a per-atom virial from it.
+    """
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_neighbor_graph,
+    )
+
+    descriptor = _build_descriptor(8)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    neighbor_graph, atype = _build_graph(descriptor, canonical=True)
+    graph = canonical_graph_from_neighbor_graph(
+        dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node)
+    )
+    fitting = _energy_fitting(descriptor)
+    operator = torch.ops.deepmd.dpa4c_canonical_compress_energy_force
+    truncated = dataclasses.replace(graph, source_row_ptr=graph.source_row_ptr[:-1])
+    arguments = _fused_arguments(
+        monkeypatch, descriptor, fitting, truncated, atype, fitting.bias_atom_e[:, 0]
+    )
+    if device == "cpu":
+        arguments = [_to_cpu(argument) for argument in arguments]
+    with pytest.raises(RuntimeError, match="source CSR"):
+        operator(*arguments)
+
+
+@_GPU
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+@pytest.mark.parametrize("spin", [False, True])
+def test_an_empty_graph_has_no_energy(
+    device: str, spin: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A graph without nodes has zero frame totals and empty node outputs.
+
+    The canonical graph of an empty system keeps its two storage slots, which
+    name a node that does not exist. The operator and its CPU reference
+    return a zero energy and virial for every frame and node outputs of the
+    declared layout without reading those slots.
+    """
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_pairs,
+    )
+
+    build = _build_spin_descriptor if spin else _build_descriptor
+    descriptor = build(8)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    empty = torch.empty(0, dtype=torch.int64, device="cuda")
+    graph = canonical_graph_from_pairs(
+        empty,
+        empty,
+        torch.empty(0, 3, dtype=torch.float32, device="cuda"),
+        torch.zeros(2, dtype=torch.int64, device="cuda"),
+        0,
+    )
+    moments = torch.empty(0, 3, dtype=torch.float32, device="cuda") if spin else None
+    fitting = _energy_fitting(descriptor)
+    arguments = _fused_arguments(
+        monkeypatch,
+        descriptor,
+        fitting,
+        graph,
+        empty,
+        fitting.bias_atom_e[:, 0],
+        moments,
+    )
+    if device == "cpu":
+        arguments = [_to_cpu(argument) for argument in arguments]
+    energy, atom_energy, force, virial, atom_virial, force_mag, mask = (
+        torch.ops.deepmd.dpa4c_canonical_compress_energy_force(*arguments)
+    )
+    assert torch.equal(energy, energy.new_zeros(2, 1))
+    assert torch.equal(virial, virial.new_zeros(2, 3, 3))
+    assert energy.dtype == virial.dtype == torch.float64
+    assert atom_energy.shape == (0, 1)
+    assert force.shape == (0, 3)
+    assert atom_virial.shape == (0, 3, 3)
+    assert force_mag.shape == ((0, 3) if spin else (0,))
+    assert mask.shape == (0,)
+
+
+@_GPU
+@pytest.mark.parametrize("frames", [1, 2])
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+@pytest.mark.parametrize("spin", [False, True])
+def test_out_of_range_forces_mark_their_frame(
+    frames: int, device: str, spin: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A force contribution of 2^30 eV/Å or more marks its frame.
+
+    The fixed-point force totals of a node hold magnitudes below 2^31 eV/Å. A
+    contribution of 2^30 eV/Å or more, the row total of a destination or an
+    edge onto its source, makes its node report its energy as not a number,
+    which marks the forces of its frame like a non-finite total; the CPU
+    reference of the operator marks the same frames. The energy is affine in
+    the head weight and the forces are linear in it, so scaling the weight
+    scales every contribution. The largest contribution of the unscaled weight
+    is measured on the CPU reference, by lowering its limit until a frame is
+    marked, which gives the scale that brings that contribution to 2^30 eV/Å: a
+    thousandth below it no frame is marked, and a thousandth above it the
+    frames that the CPU reference marks are marked, with every force of theirs
+    and none of the others. The forces scale with the weight up to the largest
+    scale that keeps every total within the range of the totals as well; the
+    check covers each contribution, and three or more contributions below the
+    limit, such as the uniform scaling makes, can carry a total out of the
+    range undetected.
+
+    A spin-conditioned model checks the magnetic contribution of every edge
+    onto its source against the same limit, in the units of the magnetic
+    force. Its moments are divided by 2^10 and its conditioning factor, the
+    gate over the reference magnitude, multiplied by 2^10: the conditioned
+    moments, and with them the energy and every coordinate contribution,
+    are bitwise those of the unscaled moments, while every magnetic
+    contribution, a cotangent of the raw moment, grows by 2^10 and reaches
+    the limit first. The frames are then marked at a far smaller weight
+    than with the unscaled moments, which the coordinate contributions
+    alone cannot cause, and the marked frames report their magnetic forces
+    as not a number as well.
+    """
+    from deepmd.pt_expt.kernels.dpa4c import (
+        canonical,
+    )
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_neighbor_graph,
+    )
+
+    build = _build_spin_descriptor if spin else _build_descriptor
+    descriptor = build(64)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    neighbor_graph, atype = _build_graph(
+        descriptor, canonical=True, node_count=40, frames=frames
+    )
+    graph = canonical_graph_from_neighbor_graph(
+        dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node)
+    )
+    fitting = _energy_fitting(descriptor)
+    operator = torch.ops.deepmd.dpa4c_canonical_compress_energy_force
+    moments = None
+    if spin:
+        generator = torch.Generator(device="cuda").manual_seed(13)
+        moments = torch.randn(
+            atype.shape[0], 3, dtype=torch.float32, device="cuda", generator=generator
+        )
+    arguments = _fused_arguments(
+        monkeypatch,
+        descriptor,
+        fitting,
+        graph,
+        atype,
+        fitting.bias_atom_e[:, 0],
+        moments,
+    )
+    reference = [_to_cpu(argument) for argument in arguments]
+
+    head_slot = [
+        argument.name for argument in operator.default._schema.arguments
+    ].index("w_head")
+
+    def evaluate(values: list[Any], scale: float) -> tuple[torch.Tensor, ...]:
+        scaled = list(values)
+        scaled[head_slot] = values[head_slot] * scale
+        return operator(*scaled)
+
+    def threshold_of(values: list[Any]) -> float:
+        """Return the weight scale that brings the largest contribution to 2^30."""
+
+        def marked_at(limit: float) -> bool:
+            with monkeypatch.context() as patch:
+                patch.setattr(canonical, "_FORCE_FIXED_LIMIT", limit)
+                energy, *_ = evaluate(values, 1.0)
+            return bool(torch.isnan(energy).any())
+
+        # The largest contribution of the unscaled weight lies between a limit
+        # that marks a frame and one that marks none; the bisection of their
+        # exponents narrows the bracket below a relative 1e-7.
+        low, high = -40.0, 30.0
+        assert marked_at(2.0**low)
+        assert not marked_at(2.0**high)
+        for _ in range(30):
+            middle = 0.5 * (low + high)
+            if marked_at(2.0**middle):
+                low = middle
+            else:
+                high = middle
+        return 2.0 ** (30.0 - high)
+
+    # === The largest contribution of the unscaled weight ===
+    threshold = threshold_of(reference)
+    if spin:
+        # === The magnetic contributions reach the limit first ===
+        names = [argument.name for argument in operator.default._schema.arguments]
+        spin_slot, type_slot = names.index("spin"), names.index("spin_type")
+
+        def rescaled(values: list[Any]) -> list[Any]:
+            values = list(values)
+            values[spin_slot] = values[spin_slot] * 2.0**-10
+            values[type_slot] = values[type_slot].clone()
+            values[type_slot][:, 0] *= 2.0**10
+            return values
+
+        arguments, reference = rescaled(arguments), rescaled(reference)
+        unscaled_threshold, threshold = threshold, threshold_of(reference)
+        assert threshold < unscaled_threshold / 16
+    tested = reference if device == "cpu" else arguments
+
+    # === Within the range, the forces scale with the weight ===
+    _, _, force, *_ = evaluate(tested, 1.0)
+    largest_total = force.abs().max().item()
+    scale = min(threshold, 2.0**31 / largest_total) * (1.0 - 1e-3)
+    energy, _, scaled, *_ = evaluate(tested, scale)
+    assert torch.isfinite(energy).all()
+    torch.testing.assert_close(
+        scaled, force * scale, rtol=1e-5, atol=1e-5 * scale * largest_total
+    )
+
+    # === Below the limit, no frame is marked ===
+    energy, *_ = evaluate(tested, threshold * (1.0 - 1e-3))
+    assert torch.isfinite(energy).all()
+
+    # === Above it, the frames of the contributions beyond it are marked ===
+    scale = threshold * (1.0 + 1e-3)
+    expected, *_ = evaluate(reference, scale)
+    energy, _, marked, _, _, marked_magnetic, _ = evaluate(tested, scale)
+    frame_marked = torch.isnan(energy).flatten().cpu()
+    assert frame_marked.any()
+    assert torch.equal(frame_marked, torch.isnan(expected).flatten())
+    node_marked = torch.repeat_interleave(frame_marked, reference[5])
+    assert node_marked.numel() == marked.shape[0]
+    marked = marked.cpu()
+    assert torch.isnan(marked[node_marked]).all()
+    assert torch.isfinite(marked[~node_marked]).all()
+    if spin:
+        marked_magnetic = marked_magnetic.cpu()
+        assert torch.isnan(marked_magnetic[node_marked]).all()
+        assert torch.isfinite(marked_magnetic[~node_marked]).all()
+
+
+def _to_cpu(value: Any) -> Any:
+    """Return a tensor, or a list of tensors, on the CPU, and anything else unchanged."""
+    if isinstance(value, torch.Tensor):
+        return value.cpu()
+    if isinstance(value, list):
+        return [_to_cpu(item) for item in value]
+    return value
+
+
+def _energy_fitting(descriptor: DescrptDPA4C) -> Any:
+    """Return the two-layer fp32 SiLU energy fitting of the fused-operator cases."""
+    from deepmd.pt_expt.fitting.ener_fitting import (
+        EnergyFittingNet,
+    )
+
+    return (
+        EnergyFittingNet(
+            ntypes=descriptor.get_ntypes(),
+            dim_descrpt=descriptor.get_dim_out(),
+            neuron=[32, 32],
+            resnet_dt=False,
+            activation_function="silu",
+            precision="float32",
+            mixed_types=True,
+            seed=29,
+        )
+        .cuda()
+        .eval()
+    )
+
+
+def _fused_arguments(
+    monkeypatch: pytest.MonkeyPatch,
+    descriptor: DescrptDPA4C,
+    fitting: Any,
+    graph: Any,
+    atype: torch.Tensor,
+    atom_bias: torch.Tensor,
+    spin: torch.Tensor | None = None,
+    pair_table: torch.Tensor | None = None,
+) -> list[Any]:
+    """Return the arguments with which the fused wrapper calls the operator.
+
+    The call is recorded rather than evaluated. The CPU reference of the
+    operator is registered first, so that the arguments can be evaluated on
+    either device.
+    """
+    from deepmd.pt_expt.kernels.dpa4c import (
+        canonical,
+    )
+
+    canonical.ensure_registered()
+    calls = []
+    with monkeypatch.context() as patch:
+        patch.setattr(
+            torch.ops.deepmd,
+            "dpa4c_canonical_compress_energy_force",
+            lambda *arguments: calls.append(list(arguments)),
+        )
+        canonical.dpa4c_canonical_compress_energy_force(
+            descriptor, fitting, graph, atype, atom_bias, spin, pair_table
+        )
+    (arguments,) = calls
+    return arguments
+
+
+def _check_canonical_against_reference(
+    descriptor: DescrptDPA4C,
+    bridged: bool,
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    node_count: int = 40,
+    side: float = 5.0,
+    single_type: bool = True,
+    halo: int = 0,
+    placeholders: bool = False,
+    frames: int = 1,
+    types: Sequence[int] = (0, 1),
+) -> tuple[torch.Tensor, ...]:
+    """Compare the fused CUDA operator with its CPU reference.
+
+    The graph holds ``frames`` frames, the last ``halo`` nodes of each left
+    unowned, and ``placeholders`` replaces the source CSR by the two-entry
+    placeholders of a caller that declines the per-atom virial. The atoms of a
+    system of several types cycle through ``types``. A bridged descriptor takes
+    the pair potential of its own type map, and is compared without it as
+    well: the tolerance follows the largest magnitude of every output, and the
+    pair potential of the short bonds of the random cluster dominates the
+    forces. A spin-conditioned descriptor receives random moments, and its
+    single-type system consists of type 0, which carries the spin branch.
+    Returns the outputs of the CUDA operator, with the pair potential of a
+    bridged descriptor.
+    """
+    from deepmd.dpmodel.atomic_model.inner_potential import (
+        InnerPotential,
+    )
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_neighbor_graph,
+    )
+
+    neighbor_graph, atype = _build_graph(
+        descriptor,
+        canonical=True,
+        node_count=node_count,
+        single_type=single_type,
+        side=side,
+        frames=frames,
+        types=types,
+    )
+    spin = None
+    if descriptor.spin is not None:
+        if single_type:
+            atype = torch.zeros_like(atype)
+        generator = torch.Generator(device="cuda").manual_seed(13)
+        spin = torch.randn(
+            atype.shape[0], 3, dtype=torch.float32, device="cuda", generator=generator
+        )
+    graph = canonical_graph_from_neighbor_graph(
+        dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node - halo)
+    )
+    if placeholders:
+        graph = dataclasses.replace(
+            graph,
+            source_row_ptr=graph.source_row_ptr.new_zeros(2),
+            source_order=graph.source_order.new_zeros(2),
+        )
+    fitting = _energy_fitting(descriptor)
+    pair_table = (
+        torch.as_tensor(
+            InnerPotential(descriptor.get_type_map()).pair_table, device="cuda"
+        )
+        if bridged
+        else None
+    )
+    operator = torch.ops.deepmd.dpa4c_canonical_compress_energy_force
+    outputs = []
+    for table in (pair_table, None) if bridged else (None,):
+        arguments = _fused_arguments(
+            monkeypatch,
+            descriptor,
+            fitting,
+            graph,
+            atype,
+            fitting.bias_atom_e[:, 0],
+            spin,
+            table,
+        )
+        kernel = operator(*arguments)
+        reference = operator(*[_to_cpu(argument) for argument in arguments])
+        for actual, expected in zip(kernel, reference, strict=True):
+            scale = max(expected.abs().max().item(), 1.0) if expected.numel() else 1.0
+            torch.testing.assert_close(
+                actual.cpu(), expected, atol=2e-5 * scale, rtol=2e-5
+            )
+        outputs.append(kernel)
+    return outputs[0]
+
+
+def _build_spin_descriptor(
+    channels: int,
+    lmax: int = 2,
+    radial_modes: int = 0,
+    device: str = "cuda",
+    type_map: Sequence[str] = tuple(_TYPE_MAP),
+    rcut: float = 3.0,
+    **window: float,
+) -> DescrptDPA4C:
+    """Return a spin-conditioned descriptor with a non-unit reference moment.
+
+    A reference magnitude other than one makes the conditioning factor visible
+    in the magnetic gradient, so a missing chain factor shows up as a constant
+    ratio rather than cancelling. The branch gate is opened, since a fresh
+    descriptor starts spin-free by design and these tests are about the branch
+    behind the gate; the gate itself is covered in ``test_dpa4c.py``. The
+    first type carries the spin branch, and the others none. The mode mixing
+    is strengthened as in `_build_descriptor`.
+    """
+    descriptor = (
+        DescrptDPA4C(
+            rcut=rcut,
+            ntypes=len(type_map),
+            channels=channels,
+            lmax=lmax,
+            n_radial=8,
+            radial_modes=radial_modes,
+            precision="float32",
+            seed=17,
+            type_map=list(type_map),
+            use_spin=[True] + [False] * (len(type_map) - 1),
+            **window,
+        )
+        .to(device)
+        .eval()
+    )
+    descriptor.spin.set_spin_reference(np.array([1.7] + [1.0] * len(type_map)))
+    with torch.no_grad():
+        descriptor.spin.spin_gate.fill_(1.0)
+    _strengthen_mode_mixing(descriptor)
+    return descriptor
+
+
+def _empty_graph(device: str) -> tuple[NeighborGraph, torch.Tensor]:
+    """Return a native graph whose node and edge axes are both empty."""
+    n_node = torch.zeros(1, dtype=torch.int64, device=device)
+    edge_index = torch.empty(2, 0, dtype=torch.int64, device=device)
+    edge_vec = torch.empty(0, 3, dtype=torch.float32, device=device)
+    edge_order = torch.empty(0, dtype=torch.int64, device=device)
+    row_pointer = torch.zeros(1, dtype=torch.int64, device=device)
+    return (
+        NeighborGraph(
+            n_node=n_node,
+            n_local=n_node,
+            edge_index=edge_index,
+            edge_vec=edge_vec,
+            edge_mask=torch.empty(0, dtype=torch.bool, device=device),
+            destination_order=edge_order,
+            destination_row_ptr=row_pointer,
+            source_order=edge_order,
+            source_row_ptr=row_pointer,
+        ),
+        torch.empty(0, dtype=torch.int64, device=device),
+    )
+
+
+def test_empty_native_spin_cpu_profile_preserves_spin_contract() -> None:
+    """The CPU reference keeps the native-spin state width at zero nodes."""
+    descriptor = _build_spin_descriptor(8, device="cpu")
+    graph, atype = _empty_graph("cpu")
+    arguments = _with_spin(
+        _arguments(descriptor, graph, atype),
+        graph.edge_vec,
+    )
+
+    output, state = _reference_forward(graph.edge_vec, *arguments)
+
+    profile = descriptor_profile(8, 2, True)
+    assert output.shape == (0, profile.output_width)
+    assert state.shape == (0, profile.state_width)
+
+
+@_GPU
+def test_empty_native_spin_cuda_backward_preserves_spin_contract() -> None:
+    """The CUDA backward distinguishes present spin from its empty axis."""
+    descriptor = _build_spin_descriptor(8)
+    graph, atype = _empty_graph("cuda")
+    arguments = _with_spin(
+        _arguments(descriptor, graph, atype),
+        graph.edge_vec,
+    )
+    output, state = torch.ops.deepmd.dpa4c_graph_compress(
+        graph.edge_vec,
+        *arguments,
+    )
+
+    edge_gradient, spin_gradient, edge_spin_gradient = _descriptor_backward(
+        torch.empty_like(output),
+        state,
+        graph.edge_vec,
+        *arguments,
+    )
+
+    assert edge_gradient.shape == (0, 3)
+    assert spin_gradient.shape == (0, 3)
+    assert edge_spin_gradient.shape == (0, 3)
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes"),
+    [(8, 2, 0), (16, 2, 0), (32, 2, 4), (64, 3, 0), (128, 3, 8), (32, 4, 0)],
+)
+@pytest.mark.parametrize("windowed", [False, True])
+@pytest.mark.parametrize("single_type", [False, True])
+def test_spin_compressed_matches_portable(
+    monkeypatch: pytest.MonkeyPatch,
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    windowed: bool,
+    single_type: bool,
+) -> None:
+    """Descriptor, coordinate gradient and magnetic force all match.
+
+    The portable path is the oracle for the spin families. Probing every
+    output column at once keeps the geometric and the spin blocks in the same
+    comparison, because the two are coupled through the shared normalizer and
+    through the cross Gram. The cotangent stays bounded: weighting columns by
+    their index inflates the gradient magnitude at the widest profile until
+    fp32 tabulation noise alone exceeds the tolerance, and the tolerances
+    follow the magnitude of each output (`_assert_close_at_scale`). A system
+    of the spin-carrying type alone is covered as well as a system of two
+    types.
+    """
+    descriptor = _build_spin_descriptor(
+        channels, lmax, radial_modes, **(_WINDOW if windowed else {})
+    )
+    assert mega_eligible(descriptor)
+    graph, atype = _build_graph(descriptor, canonical=True)
+    if single_type:
+        # Type 0 carries the spin branch; the graph does not depend on types.
+        atype = torch.zeros_like(atype)
+    generator = torch.Generator(device="cuda").manual_seed(11)
+    spin = torch.randn(
+        atype.shape[0],
+        3,
+        dtype=torch.float32,
+        device="cuda",
+        generator=generator,
     )
 
     def run() -> tuple[torch.Tensor, ...]:
@@ -1624,212 +3243,1558 @@ def run() -> tuple[torch.Tensor, ...]:
             atype,
             spin=moment,
         )
-        cotangent = torch.linspace(
-            -0.7,
-            1.3,
-            output.numel(),
-            dtype=output.dtype,
-            device=output.device,
-        ).reshape_as(output)
-        gradients = torch.autograd.grad(
-            (output * cotangent).sum(),
-            [edge_vec, moment],
-        )
-        return (output, *gradients)
+        cotangent = torch.linspace(
+            -0.7,
+            1.3,
+            output.numel(),
+            dtype=output.dtype,
+            device=output.device,
+        ).reshape_as(output)
+        gradients = torch.autograd.grad(
+            (output * cotangent).sum(),
+            [edge_vec, moment],
+        )
+        return (output, *gradients)
+
+    reference = run()
+    descriptor.enable_compression(0.5)
+    monkeypatch.setenv("DP_CUDA_INFER", "1")
+    assert descriptor.compress
+    actual = run()
+    _assert_dispatched(actual[0], reference[0])
+    # Tabulation error reaches the gradients through the table derivative, so
+    # they carry the wider tolerance the geometric backward already uses, and
+    # the steep switch inside a bridging window widens it as in the surface
+    # test.
+    gradient = (1e-4, 5e-4) if windowed else (8e-6, 1e-4)
+    tolerances = ((3e-5, 3e-5), gradient, gradient)
+    for (atol, rtol), value, expected in zip(
+        tolerances, actual, reference, strict=True
+    ):
+        _assert_close_at_scale(value, expected, atol=atol, rtol=rtol)
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "type_count", "windowed"),
+    [
+        (64, 3, 4, 1, False),
+        (64, 3, 4, 3, False),
+        (64, 3, 4, 3, True),
+        (128, 3, 4, 1, False),
+        (128, 3, 4, 4, False),
+        (128, 3, 4, 4, True),
+        (32, 2, 4, 2, False),
+        (64, 3, 2, 3, False),
+        (32, 2, 2, 2, False),
+        (128, 3, 2, 1, False),
+        (64, 3, 4, 6, False),
+        (64, 3, 4, 8, True),
+        (128, 3, 4, 6, False),
+        (128, 3, 4, 8, True),
+        (64, 3, 2, 6, False),
+        (128, 3, 8, 2, False),
+        (128, 3, 8, 2, True),
+    ],
+)
+@pytest.mark.parametrize("system", ["unsplit", "coscheduled"])
+def test_spin_slotted_scans_match_the_portable_equations(
+    monkeypatch: pytest.MonkeyPatch,
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    type_count: int,
+    windowed: bool,
+    system: str,
+) -> None:
+    """The slotted spin scans follow the portable equations.
+
+    A canonical spin system of one to eight node types slots the mode profiles
+    of its edges by the types of their sources in the edge scans of both
+    directions, in four slots up to four types and in eight from five, for two
+    modes as well as for four. The graph carries 64-bit indices, whose rows
+    always take a warp each; a co-scheduled system holds four waves of single
+    24-row blocks per multiprocessor, more than the co-scheduled blocks of
+    every slotted spin scan need. A descriptor of eight radial modes, more than
+    the slots carry, scans pairwise, and the co-scheduled system then in the
+    co-scheduled blocks of the pairwise spin scans. The atoms of the random
+    cluster keep the density of the 24-atom one. Descriptor, coordinate
+    gradient and magnetic force follow the portable path, at tolerances that
+    follow the magnitude of each output (`_assert_close_at_scale`).
+    """
+    multiprocessors = torch.cuda.get_device_properties("cuda").multi_processor_count
+    node_count = {
+        "unsplit": 2 * multiprocessors + 17,
+        "coscheduled": 4 * 24 * multiprocessors + 17,
+    }[system]
+    descriptor = _build_spin_descriptor(
+        channels,
+        lmax,
+        radial_modes,
+        type_map=_WIDE_TYPE_MAP[: max(type_count, 4)],
+        **(_WINDOW if windowed else {}),
+    )
+    graph, atype = _build_graph(
+        descriptor,
+        canonical=True,
+        node_count=node_count,
+        side=5.0 * (node_count / 24) ** (1.0 / 3.0),
+        types=range(type_count),
+    )
+    generator = torch.Generator(device="cuda").manual_seed(11)
+    spin = torch.randn(
+        atype.shape[0], 3, dtype=torch.float32, device="cuda", generator=generator
+    )
+
+    def run() -> tuple[torch.Tensor, ...]:
+        moment = spin.detach().clone().requires_grad_(True)
+        edge_vec = graph.edge_vec.detach().clone().requires_grad_(True)
+        output, _ = descriptor.call_graph(
+            dataclasses.replace(graph, edge_vec=edge_vec), atype, spin=moment
+        )
+        cotangent = torch.linspace(
+            -0.7, 1.3, output.numel(), dtype=output.dtype, device=output.device
+        ).reshape_as(output)
+        gradients = torch.autograd.grad((output * cotangent).sum(), [edge_vec, moment])
+        return (output, *gradients)
+
+    reference = run()
+    descriptor.enable_compression(0.5)
+    monkeypatch.setenv("DP_CUDA_INFER", "1")
+    actual = run()
+    _assert_dispatched(actual[0], reference[0])
+    gradient = (1e-4, 5e-4) if windowed else (8e-6, 1e-4)
+    tolerances = ((3e-5, 3e-5), gradient, gradient)
+    for (atol, rtol), value, expected in zip(
+        tolerances, actual, reference, strict=True
+    ):
+        _assert_close_at_scale(value, expected, atol=atol, rtol=rtol)
+
+
+@_GPU
+@pytest.mark.parametrize(
+    ("channels", "lmax", "radial_modes", "slots"),
+    [(64, 3, 4, True), (128, 3, 4, True), (64, 2, 0, False)],
+)
+@pytest.mark.parametrize("system", ["unsplit", "coscheduled"])
+@pytest.mark.parametrize("slot_types", [1, 3, 6])
+def test_spin_systems_of_a_few_types_take_the_slotted_scans(
+    channels: int,
+    lmax: int,
+    radial_modes: int,
+    slots: bool,
+    system: str,
+    slot_types: int,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A spin system of one to eight node types takes the slotted scans.
+
+    Both systems hold the same atoms of the spin-carrying type, and eight
+    isolated atoms, which interact with none: in one the isolated atoms carry
+    that type, or two or five other types besides it, a system of one, three or
+    six types, and in the other the eight other types of the descriptor, a
+    system of nine types that the slotted scans do not serve. The node counts
+    agree, so every launch choice that follows from them does as well, and the
+    rows of both take a warp each. The slotted scans apply the mode mixing of
+    each slot to reduced sums, so their forces and magnetic forces differ from
+    the pairwise ones in the last bits. A descriptor without radial modes has
+    nothing to slot and scans pairwise either way: its outputs agree bit for
+    bit. The forces of the system of one type follow the CPU reference of the
+    fused operator.
+    """
+    from deepmd.pt_expt.kernels.dpa4c.canonical import (
+        dpa4c_canonical_compress_energy_force,
+    )
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_neighbor_graph,
+    )
+
+    multiprocessors = torch.cuda.get_device_properties("cuda").multi_processor_count
+    node_count = {
+        "unsplit": 16 * multiprocessors + 17,
+        "coscheduled": 4 * 24 * multiprocessors + 17,
+    }[system]
+    descriptor = _build_spin_descriptor(
+        channels, lmax, radial_modes, type_map=_WIDE_TYPE_MAP
+    )
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    side = 5.0 * (node_count / 40) ** (1.0 / 3.0)
+    if slot_types == 1:
+        _check_canonical_against_reference(
+            descriptor, False, monkeypatch, node_count=node_count, side=side
+        )
+    neighbor_graph, atype = _build_graph(
+        descriptor, canonical=True, node_count=node_count, side=side
+    )
+    graph = canonical_graph_from_neighbor_graph(
+        dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node)
+    )
+    atype = torch.zeros_like(atype)
+    # The isolated atoms close both row-pointer arrays with empty rows.
+    isolated = 8
+    extended = dataclasses.replace(
+        graph,
+        n_node=graph.n_node + isolated,
+        n_local=graph.n_local + isolated,
+        destination_row_ptr=torch.cat(
+            [graph.destination_row_ptr] + [graph.destination_row_ptr[-1:]] * isolated
+        ),
+        source_row_ptr=torch.cat(
+            [graph.source_row_ptr] + [graph.source_row_ptr[-1:]] * isolated
+        ),
+    )
+    isolated_types = {
+        1: [0, 0, 0, 0, 0, 0, 0, 0],
+        3: [1, 2, 0, 0, 0, 0, 0, 0],
+        6: [1, 2, 3, 4, 5, 0, 0, 0],
+    }[slot_types]
+    few_types = torch.cat([atype, atype.new_tensor(isolated_types)])
+    all_types = torch.cat([atype, atype.new_tensor([1, 2, 3, 4, 5, 6, 7, 8])])
+    generator = torch.Generator(device="cuda").manual_seed(13)
+    moments = torch.randn(
+        few_types.shape[0], 3, dtype=torch.float32, device="cuda", generator=generator
+    )
+    fitting = _energy_fitting(descriptor)
+    nodes = atype.shape[0]
+    slotted = dpa4c_canonical_compress_energy_force(
+        descriptor, fitting, extended, few_types, fitting.bias_atom_e[:, 0], moments
+    )
+    pairwise = dpa4c_canonical_compress_energy_force(
+        descriptor, fitting, extended, all_types, fitting.bias_atom_e[:, 0], moments
+    )
+    scale = max(pairwise[2].abs().max().item(), 1.0)
+    torch.testing.assert_close(
+        slotted[2][:nodes], pairwise[2][:nodes], atol=2e-5 * scale, rtol=2e-5
+    )
+    torch.testing.assert_close(
+        slotted[5][:nodes], pairwise[5][:nodes], atol=2e-5, rtol=2e-5
+    )
+    assert torch.equal(slotted[2][:nodes], pairwise[2][:nodes]) is not slots
+    assert torch.equal(slotted[5][:nodes], pairwise[5][:nodes]) is not slots
+
+
+@_GPU
+@pytest.mark.parametrize("family", ["neighbour", "onsite"])
+def test_spin_magnetic_force_splits_into_onsite_and_neighbour(
+    monkeypatch: pytest.MonkeyPatch,
+    family: str,
+) -> None:
+    """Each half of the magnetic force is correct on its own.
+
+    The on-site half closes inside the node kernel while the neighbour half is
+    emitted per edge and reduced onto source nodes, so they fail
+    independently. Silencing one at a time keeps a fault in either from being
+    masked by the other's magnitude.
+    """
+    descriptor = _build_spin_descriptor(16)
+    with torch.no_grad():
+        if family == "neighbour":
+            descriptor.spin.adam_spin_vector_weight.zero_()
+            descriptor.spin.adam_spin_quadrupole_weight.zero_()
+        else:
+            geometric = descriptor.channels * (2 + descriptor.radial_modes)
+            for parameter in descriptor.pair_film.parameters():
+                if parameter.dim() == 2 and parameter.shape[1] > geometric:
+                    parameter[:, geometric:] = 0.0
+                elif parameter.dim() == 1 and parameter.shape[0] > geometric:
+                    parameter[geometric:] = 0.0
+    graph, atype = _build_graph(descriptor, canonical=True)
+    generator = torch.Generator(device="cuda").manual_seed(11)
+    spin = torch.randn(
+        atype.shape[0],
+        3,
+        dtype=torch.float32,
+        device="cuda",
+        generator=generator,
+    )
+
+    def run() -> tuple[torch.Tensor, torch.Tensor]:
+        moment = spin.detach().clone().requires_grad_(True)
+        output, _ = descriptor.call_graph(graph, atype, spin=moment)
+        cotangent = torch.linspace(
+            -0.7,
+            1.3,
+            output.numel(),
+            dtype=output.dtype,
+            device=output.device,
+        ).reshape_as(output)
+        gradient = torch.autograd.grad((output * cotangent).sum(), moment)[0]
+        return output, gradient
+
+    reference_output, reference_gradient = run()
+    descriptor.enable_compression(0.5)
+    monkeypatch.setenv("DP_CUDA_INFER", "1")
+    output, gradient = run()
+    _assert_dispatched(output, reference_output)
+    assert torch.count_nonzero(gradient).item() > 0
+    torch.testing.assert_close(gradient, reference_gradient, atol=8e-6, rtol=1e-4)
+
+
+@_GPU
+def test_spin_bond_family_couples_the_kernel_to_the_edge_direction(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The compiled spin path reads the edge direction, and only through it.
+
+    Rotating the neighbourhood while holding the moments fixed leaves every
+    spin family except the bond-projected one invariant, so a kernel that
+    omitted that family would return an unchanged readout. Rotating the
+    moments alongside the geometry restores the invariance, which is what
+    separates a genuine bond coupling from a defect in the geometric block.
+    """
+    descriptor = _build_spin_descriptor(16)
+    graph, atype = _build_graph(descriptor, canonical=True)
+    generator = torch.Generator(device="cuda").manual_seed(11)
+    spin = torch.randn(
+        atype.shape[0],
+        3,
+        dtype=torch.float32,
+        device="cuda",
+        generator=generator,
+    )
+    angle = 0.4
+    rotation = torch.tensor(
+        [
+            [float(np.cos(angle)), -float(np.sin(angle)), 0.0],
+            [float(np.sin(angle)), float(np.cos(angle)), 0.0],
+            [0.0, 0.0, 1.0],
+        ],
+        dtype=torch.float32,
+        device="cuda",
+    )
+
+    def run(edge_vec: torch.Tensor, moment: torch.Tensor) -> torch.Tensor:
+        output, _ = descriptor.call_graph(
+            dataclasses.replace(graph, edge_vec=edge_vec),
+            atype,
+            spin=moment,
+        )
+        return output
 
-    reference = run()
+    portable = run(graph.edge_vec, spin)
     descriptor.enable_compression(0.5)
     monkeypatch.setenv("DP_CUDA_INFER", "1")
-    assert descriptor.compress
-    actual = run()
-    _assert_dispatched(actual[0], reference[0])
-    # Tabulation error reaches the gradients through the table derivative, so
-    # they carry the wider tolerance the geometric backward already uses, and
-    # the steep switch inside a bridging window widens it as in the surface
-    # test.
-    gradient = (1e-4, 5e-4) if windowed else (8e-6, 1e-4)
-    tolerances = ((3e-5, 3e-5), gradient, gradient)
-    for (atol, rtol), value, expected in zip(
-        tolerances, actual, reference, strict=True
+    upright = run(graph.edge_vec, spin)
+    _assert_dispatched(upright, portable)
+    rotated = run(graph.edge_vec @ rotation.T, spin)
+    assert not torch.allclose(rotated, upright, atol=1e-4)
+    covariant = run(graph.edge_vec @ rotation.T, spin @ rotation.T)
+    torch.testing.assert_close(covariant, upright, atol=3e-5, rtol=3e-5)
+
+
+@_GPU
+@pytest.mark.parametrize("bridged", [False, True])
+@pytest.mark.parametrize("spin_conditioned", [False, True])
+def test_backward_operator_satisfies_its_schema(
+    spin_conditioned: bool,
+    bridged: bool,
+) -> None:
+    """The backward operator declares four independent results.
+
+    All four are unannotated, so any two of them sharing storage would be an
+    alias the schema does not describe, which is undefined under
+    functionalization. ``opcheck`` decides that mechanically, including on the
+    spin-free path without a pair potential, where three of the four are
+    absent and an allocation shared between them would otherwise go unnoticed.
+    """
+    from deepmd.dpmodel.atomic_model.inner_potential import (
+        InnerPotential,
+    )
+
+    descriptor = _build_spin_descriptor(8) if spin_conditioned else _build_descriptor(8)
+    graph, atype = _build_graph(descriptor, canonical=True)
+    arguments = _arguments(descriptor, graph, atype)
+    pair = (
+        (
+            torch.as_tensor(InnerPotential(["O", "H"]).pair_table, device="cuda"),
+            torch.ones(atype.shape[0], dtype=torch.float64, device="cuda"),
+        )
+        if bridged
+        else (graph.edge_vec.new_empty(0),) * 2
+    )
+    if spin_conditioned:
+        arguments = _with_spin(
+            arguments,
+            torch.randn(atype.shape[0], 3, dtype=torch.float32, device="cuda"),
+        )
+    output, state = torch.ops.deepmd.dpa4c_graph_compress(
+        graph.edge_vec,
+        *arguments,
+    )
+    torch.library.opcheck(
+        torch.ops.deepmd.dpa4c_graph_compress_backward.default,
+        (torch.ones_like(output), state, graph.edge_vec, *arguments, *pair),
+    )
+
+
+@_GPU
+def test_registered_autograd_refuses_the_magnetic_moment() -> None:
+    """A direct operator call cannot silently lose the magnetic force.
+
+    The operator emits that cotangent in two pieces and the per-edge piece is
+    reduced onto source nodes through the source CSR, which the schema does
+    not carry. The registration therefore cannot close the magnetic force, and
+    refusing is the only alternative to reporting a vanishing one.
+    """
+    descriptor = _build_spin_descriptor(8)
+    graph, atype = _build_graph(descriptor, canonical=True)
+    spin = torch.randn(
+        atype.shape[0],
+        3,
+        dtype=torch.float32,
+        device="cuda",
+        requires_grad=True,
+    )
+    output, _state = torch.ops.deepmd.dpa4c_graph_compress(
+        graph.edge_vec,
+        *_with_spin(_arguments(descriptor, graph, atype), spin),
+    )
+    with pytest.raises(RuntimeError, match="cannot differentiate its magnetic"):
+        output.sum().backward()
+
+
+# === Fitting routes of the fused energy-force operator ===
+
+
+def _route_fitting(descriptor: DescrptDPA4C, width: int) -> Any:
+    """Return a three-layer fp32 SiLU energy fitting of hidden width ``width``."""
+    from deepmd.pt_expt.fitting.ener_fitting import (
+        EnergyFittingNet,
+    )
+
+    return (
+        EnergyFittingNet(
+            ntypes=descriptor.get_ntypes(),
+            dim_descrpt=descriptor.get_dim_out(),
+            neuron=[width] * 3,
+            resnet_dt=False,
+            activation_function="silu",
+            precision="float32",
+            mixed_types=True,
+            seed=31,
+        )
+        .cuda()
+        .eval()
+    )
+
+
+def _route_case(node_count: int) -> tuple[DescrptDPA4C, Any, torch.Tensor]:
+    """Return a compressed descriptor and the canonical graph of ``node_count`` atoms."""
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_neighbor_graph,
+    )
+
+    descriptor = _build_descriptor(8, 2, 0)
+    descriptor._set_compression(build_compression_artifacts(descriptor))
+    neighbor_graph, atype = _build_graph(
+        descriptor,
+        canonical=True,
+        node_count=node_count,
+        side=5.0 * (node_count / 40) ** (1.0 / 3.0),
+    )
+    graph = canonical_graph_from_neighbor_graph(
+        dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node)
+    )
+    return descriptor, graph, atype
+
+
+def _single_route(choice: int) -> list[int]:
+    """Return the fitting-route table that names ``choice`` for every run size."""
+    from deepmd.pt_expt.kernels.fitting_routes import (
+        ROUTE_TABLE_MAGIC,
+    )
+
+    fingerprint = torch.ops.deepmd.graph_fitting_route_fingerprint(
+        torch.cuda.current_device()
+    )
+    return [ROUTE_TABLE_MAGIC, *fingerprint, 1, choice]
+
+
+def _fitting_kernels(call: Callable[[], Any]) -> list[str]:
+    """Return the kinds of the fitting kernels that ``call`` launches, sorted.
+
+    A trace that records no CUDA kernel at all fails the check rather than
+    reading as a route without fitting kernels.
+    """
+    from torch.profiler import (
+        ProfilerActivity,
+        profile,
+    )
+
+    with profile(activities=[ProfilerActivity.CUDA]) as trace:
+        call()
+        torch.cuda.synchronize()
+    names = [
+        event.name
+        for event in trace.events()
+        if event.device_type == torch.autograd.DeviceType.CUDA
+    ]
+    assert names, "the profiler recorded no CUDA kernel"
+    kinds = set()
+    for name in names:
+        if "fused_network_kernel" in name:
+            kinds.add("fused")
+        elif "layer_forward_kernel" in name:
+            kinds.add("per-layer")
+    return sorted(kinds)
+
+
+#: Run sizes of the route cases, in multiprocessors of the device: a run of
+#: fewer 32-node panels than multiprocessors and one of more.
+_ROUTE_RUNS = {"latency": 0.5, "throughput": 3.5}
+
+
+@_GPU
+@pytest.mark.parametrize("width", [64, 96, 192, 256, 384])  # fused hidden width
+@pytest.mark.parametrize("run", sorted(_ROUTE_RUNS))  # run size regime
+def test_every_fitting_route_returns_the_bits_of_the_rules(
+    width: int, run: str
+) -> None:
+    """Every route of the fitting network gives the outputs of the rules bit for bit.
+
+    A route table that names the per-layer kernels or a fused shape the device
+    runs for the run size routes the network there, and every route evaluates
+    the same contractions in the same order, so all seven outputs of the
+    operator equal those of the rules, which an empty table takes.
+    """
+    from deepmd.pt_expt.kernels.dpa4c.canonical import (
+        dpa4c_canonical_compress_energy_force,
+    )
+    from deepmd.pt_expt.kernels.fitting_routes import (
+        PER_LAYER,
+    )
+
+    multiprocessors = torch.cuda.get_device_properties("cuda").multi_processor_count
+    node_count = int(32 * multiprocessors * _ROUTE_RUNS[run])
+    descriptor, graph, atype = _route_case(node_count)
+    fitting = _route_fitting(descriptor, width)
+    bias = fitting.bias_atom_e[:, 0]
+    rules = dpa4c_canonical_compress_energy_force(
+        descriptor, fitting, graph, atype, bias, fitting_routes=[]
+    )
+    choices = [
+        PER_LAYER,
+        *torch.ops.deepmd.graph_fitting_route_choices(
+            torch.cuda.current_device(), width, node_count
+        ),
+    ]
+    for choice in choices:
+        routed = dpa4c_canonical_compress_energy_force(
+            descriptor,
+            fitting,
+            graph,
+            atype,
+            bias,
+            fitting_routes=_single_route(choice),
+        )
+        for name, expected, actual in zip(
+            (
+                "energy",
+                "atom_energy",
+                "force",
+                "virial",
+                "atom_virial",
+                "force_mag",
+                "mask",
+            ),
+            rules,
+            routed,
+            strict=True,
+        ):
+            assert torch.equal(expected, actual), (choice, name)
+
+
+@_GPU
+@pytest.mark.parametrize("width", [96, 192])  # fused hidden width
+def test_a_tiled_call_routes_its_runs_and_its_tail_apart(
+    width: int, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """The full runs and the tail of a tiled call take their routes apart.
+
+    The operator routes every run of a tiled call by the run's own size, so a
+    table that splits the run sizes at 1,000 nodes sends full runs of 3.5
+    waves of panels and a tail of 100 nodes to the per-layer kernels, the
+    rules or a fused shape in each combination, and the per-layer buffers
+    serve whichever runs take those kernels. Every combination gives the
+    outputs of the untiled call under the rules bit for bit.
+    """
+    from deepmd.pt_expt.kernels.dpa4c.canonical import (
+        dpa4c_canonical_compress_energy_force,
+    )
+    from deepmd.pt_expt.kernels.fitting_routes import (
+        PER_LAYER,
+        ROUTE_TABLE_MAGIC,
+        RULES,
+    )
+    from deepmd.pt_expt.kernels.graph_fitting import (
+        fitting_operator_arguments,
+    )
+
+    device = torch.cuda.current_device()
+    multiprocessors = torch.cuda.get_device_properties(device).multi_processor_count
+    tile, tail = int(32 * multiprocessors * _ROUTE_RUNS["throughput"]), 100
+    descriptor, graph, atype = _route_case(2 * tile + tail)
+    fitting = _route_fitting(descriptor, width)
+    bias = fitting.bias_atom_e[:, 0]
+    weights = list(fitting_operator_arguments(fitting).weights)
+    fingerprint = torch.ops.deepmd.graph_fitting_route_fingerprint(device)
+
+    def run(table: list[int]) -> tuple[torch.Tensor, ...]:
+        return dpa4c_canonical_compress_energy_force(
+            descriptor, fitting, graph, atype, bias, fitting_routes=table
+        )
+
+    monkeypatch.setenv("DP_NODE_TILE", "0")
+    reference = run([])
+    monkeypatch.setenv("DP_NODE_TILE", str(tile))
+    combinations = [(PER_LAYER, RULES), (RULES, PER_LAYER), (PER_LAYER, PER_LAYER)]
+    tail_shapes = torch.ops.deepmd.graph_fitting_route_choices(device, width, tail)
+    full_shapes = torch.ops.deepmd.graph_fitting_route_choices(device, width, tile)
+    if tail_shapes and full_shapes:
+        combinations.append((tail_shapes[-1], full_shapes[0]))
+    for tail_route, full_route in combinations:
+        table = [ROUTE_TABLE_MAGIC, *fingerprint, 2, 1000, tail_route, full_route]
+        for size, route in ((tail, tail_route), (tile, full_route)):
+            if route != RULES:
+                resolved = torch.ops.deepmd.graph_fitting_route_resolve(
+                    weights, size, table
+                )
+                assert resolved == route, (size, route)
+        for actual, expected in zip(run(table), reference, strict=True):
+            assert torch.equal(actual, expected), (tail_route, full_route)
+
+
+def _route_routing_case() -> tuple[DescrptDPA4C, Any, Any, torch.Tensor]:
+    """Return the network of width 192 and a run of four waves of its panels."""
+    multiprocessors = torch.cuda.get_device_properties("cuda").multi_processor_count
+    descriptor, graph, atype = _route_case(32 * multiprocessors * 4)
+    return descriptor, _route_fitting(descriptor, 192), graph, atype
+
+
+def _route_tables() -> dict[str, list[int]]:
+    """Return the tables of the routing cases, by name."""
+    from deepmd.pt_expt.kernels.fitting_routes import (
+        PER_LAYER,
+        RULES,
+    )
+
+    per_layer = _single_route(PER_LAYER)
+    foreign = list(per_layer)
+    foreign[2] += 1  # the multiprocessor count of the fingerprint
+    return {
+        "empty": [],
+        "rules": _single_route(RULES),
+        "per-layer": per_layer,
+        "foreign": foreign,
+        "malformed": per_layer[:-1],
+        "unknown shape": _single_route(0x7FFF_FFFF_FFFF),
+    }
+
+
+def _traced_routes() -> dict[str, list[str]]:
+    """Return the fitting kernels that each routing table launches, by name."""
+    from deepmd.pt_expt.kernels.dpa4c.canonical import (
+        dpa4c_canonical_compress_energy_force,
+    )
+
+    descriptor, fitting, graph, atype = _route_routing_case()
+    bias = fitting.bias_atom_e[:, 0]
+    return {
+        name: _fitting_kernels(
+            lambda table=table: dpa4c_canonical_compress_energy_force(
+                descriptor, fitting, graph, atype, bias, fitting_routes=table
+            )
+        )
+        for name, table in _route_tables().items()
+    }
+
+
+@_GPU
+def test_a_fitting_route_table_routes_the_network() -> None:
+    """The table decides the route; an unusable table leaves it to the rules.
+
+    At a width of 192 the rules route a run of a few waves of panels to the
+    fused kernel. A table that names the per-layer kernels routes it there; a
+    table that names the rules, one measured on a device of another
+    fingerprint, a malformed one and one that names a shape not compiled for
+    the width leave the run to the rules.
+    The routes resolve through the query of the operator, and the kernels they
+    launch, which a profiler trace in a process of its own shows, follow them.
+    """
+    from deepmd.pt_expt.kernels.fitting_routes import (
+        PER_LAYER,
+    )
+    from deepmd.pt_expt.kernels.graph_fitting import (
+        fitting_operator_arguments,
+    )
+
+    _descriptor, fitting, graph, _atype = _route_routing_case()
+    weights = list(fitting_operator_arguments(fitting).weights)
+    nodes = graph.destination_row_ptr.numel() - 1
+    resolved = {
+        name: torch.ops.deepmd.graph_fitting_route_resolve(weights, nodes, table)
+        for name, table in _route_tables().items()
+    }
+    assert resolved == {
+        "empty": 0,
+        "rules": 0,
+        "per-layer": PER_LAYER,
+        "foreign": 0,
+        "malformed": 0,
+        "unknown shape": 0,
+    }
+    traced = _run_in_process("_traced_routes()")
+    assert traced == {
+        "empty": ["fused"],
+        "rules": ["fused"],
+        "per-layer": ["per-layer"],
+        "foreign": ["fused"],
+        "malformed": ["fused"],
+        "unknown shape": ["fused"],
+    }
+
+
+def _route_model(compressed: bool = True) -> Any:
+    """Return a DPA4C energy model whose fitting has a hidden width of 192."""
+    from deepmd.pt_expt.model.get_model import (
+        get_model,
+    )
+
+    model = (
+        get_model(
+            {
+                "type_map": ["O", "H"],
+                "descriptor": {
+                    "type": "dpa4c",
+                    "rcut": 4.0,
+                    "channels": 8,
+                    "lmax": 2,
+                    "n_radial": 8,
+                    "precision": "float32",
+                    "seed": 3,
+                },
+                "fitting_net": {
+                    "type": "ener",
+                    "neuron": [192, 192, 192],
+                    "activation_function": "silu",
+                    "precision": "float32",
+                    "resnet_dt": False,
+                    "seed": 3,
+                },
+            }
+        )
+        .to("cpu")
+        .eval()
+    )
+    if compressed:
+        model.get_descriptor().enable_compression(min_nbor_dist=0.5)
+    return model
+
+
+def _frozen_route_kernels(directory: str) -> dict[str, Any]:
+    """Freeze the routing model under two held tables; return what each archive does.
+
+    The tuner skips a network whose table it holds, so a table set beforehand
+    stands for a measured one. Each archive is frozen to ``directory`` and run
+    through AOTInductor, as a deployment runs it. Returns, for the table that
+    names the per-layer kernels and for no table, the table argument of the
+    fused operator in the exported program and the fitting kernels that the
+    archive launches.
+    """
+    from torch._inductor import (
+        aoti_load_package,
+    )
+
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_index_bytes,
+    )
+    from deepmd.pt_expt.utils.serialization import (
+        _trace_and_export,
+        deserialize_to_file,
+    )
+
+    os.environ["DP_CUDA_INFER"] = "2"
+    model = _route_model()
+    learned = model.atomic_model.fused_decomposition()[0]
+    key = fitting_routes.network_key(learned.fitting_net)
+    device = torch.device("cuda", torch.cuda.current_device())
+    fingerprint = fitting_routes._fingerprint(device)
+    multiprocessors = torch.cuda.get_device_properties(device).multi_processor_count
+    node_count = 32 * multiprocessors * 4
+    graph = fitting_routes._bulk_graph(4.0, node_count, device)
+    atype = torch.arange(node_count, device=device) % 2
+    inputs = (
+        atype,
+        graph.n_node,
+        graph.n_local,
+        canonical_index_bytes(graph.source),
+        graph.edge_vec,
+        graph.destination_row_ptr,
+        graph.source_row_ptr,
+        graph.source_order,
+    )
+    results = {}
+    for name, table in (
+        ("per-layer", _single_route(fitting_routes.PER_LAYER)),
+        ("rules", []),
     ):
-        torch.testing.assert_close(value, expected, atol=atol, rtol=rtol)
+        fitting_routes._ROUTE_TABLES[(fingerprint, key)] = tuple(table)
+        exported, _metadata, _model_json, _output_keys = _trace_and_export(
+            {"model": model.serialize()},
+            lower_kind="dpa4c_canonical",
+            do_atomic_virial=True,
+        )
+        # The table is the last argument of the schema; a call that leaves it
+        # to its default of no table records one argument fewer.
+        schema = torch.ops.deepmd.dpa4c_canonical_compress_energy_force.default._schema
+        literals = [
+            node.args[-1] if len(node.args) == len(schema.arguments) else None
+            for node in exported.graph_module.graph.nodes
+            if "dpa4c_canonical_compress_energy_force" in str(node.target)
+        ]
+        archive = os.path.join(directory, f"{name}.pt2")
+        deserialize_to_file(
+            archive,
+            {"model": model.serialize()},
+            do_atomic_virial=True,
+            lower_kind="dpa4c_canonical",
+            with_comm=False,
+        )
+        runner = aoti_load_package(archive)
+        results[name] = {
+            "literals": [None if v is None else [int(x) for x in v] for v in literals],
+            "kernels": _fitting_kernels(lambda runner=runner: runner(*inputs)),
+        }
+    return results
 
 
 @_GPU
-@pytest.mark.parametrize("family", ["neighbour", "onsite"])
-def test_spin_magnetic_force_splits_into_onsite_and_neighbour(
+def test_the_measured_routes_reach_the_frozen_program(tmp_path: pathlib.Path) -> None:
+    """The table held for a network routes the frozen archive of the model.
+
+    The export of the compact canonical lower takes the held table as an
+    argument of the fused operator, or no table where none is held, and the
+    frozen archive, run through AOTInductor, routes the network as the table
+    names: through the per-layer kernels under a table that names them, and
+    through the fused kernel, which the rules choose for the run, without a
+    table. The archives run in a process of their own, whose profiler trace
+    shows the kernels.
+    """
+    from deepmd.pt_expt.kernels.fitting_routes import (
+        PER_LAYER,
+    )
+
+    results = _run_in_process(f"_frozen_route_kernels({str(tmp_path)!r})")
+    assert results == {
+        "per-layer": {
+            "literals": [_single_route(PER_LAYER)],
+            "kernels": ["per-layer"],
+        },
+        "rules": {"literals": [None], "kernels": ["fused"]},
+    }
+
+
+@_GPU
+def test_the_route_tuner_leaves_other_models_alone(
     monkeypatch: pytest.MonkeyPatch,
-    family: str,
 ) -> None:
-    """Each half of the magnetic force is correct on its own.
+    """The tuner measures nothing for a model off the fused DPA4C operator.
 
-    The on-site half closes inside the node kernel while the neighbour half is
-    emitted per edge and reduced onto source nodes, so they fail
-    independently. Silencing one at a time keeps a fault in either from being
-    masked by the other's magnitude.
+    Nor does a freeze of a model that takes the operator measure anything for
+    a lower that does not call it.
     """
-    descriptor = _build_spin_descriptor(16)
-    with torch.no_grad():
-        if family == "neighbour":
-            descriptor.spin.adam_spin_vector_weight.zero_()
-            descriptor.spin.adam_spin_quadrupole_weight.zero_()
-        else:
-            geometric = descriptor.channels * (2 + descriptor.radial_modes)
-            for parameter in descriptor.pair_film.parameters():
-                if parameter.dim() == 2 and parameter.shape[1] > geometric:
-                    parameter[:, geometric:] = 0.0
-                elif parameter.dim() == 1 and parameter.shape[0] > geometric:
-                    parameter[geometric:] = 0.0
-    graph, atype = _build_graph(descriptor, canonical=True)
-    generator = torch.Generator(device="cuda").manual_seed(11)
-    spin = torch.randn(
-        atype.shape[0],
-        3,
-        dtype=torch.float32,
-        device="cuda",
-        generator=generator,
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+    from deepmd.pt_expt.kernels.autotune import (
+        run_autotune,
     )
 
-    def run() -> tuple[torch.Tensor, torch.Tensor]:
-        moment = spin.detach().clone().requires_grad_(True)
-        output, _ = descriptor.call_graph(graph, atype, spin=moment)
-        cotangent = torch.linspace(
-            -0.7,
-            1.3,
-            output.numel(),
-            dtype=output.dtype,
-            device=output.device,
-        ).reshape_as(output)
-        gradient = torch.autograd.grad((output * cotangent).sum(), moment)[0]
-        return output, gradient
-
-    reference_output, reference_gradient = run()
-    descriptor.enable_compression(0.5)
-    monkeypatch.setenv("DP_CUDA_INFER", "1")
-    output, gradient = run()
-    _assert_dispatched(output, reference_output)
-    assert torch.count_nonzero(gradient).item() > 0
-    torch.testing.assert_close(gradient, reference_gradient, atol=8e-6, rtol=1e-4)
+    tables: dict = {}
+    monkeypatch.setattr(fitting_routes, "_ROUTE_TABLES", tables)
+    device = torch.device("cuda", torch.cuda.current_device())
+    for model in (_route_model(compressed=False), torch.nn.Linear(2, 2)):
+        fitting_routes.tune_fitting_routes(model, 0, device)
+    run_autotune(_route_model(), device, "graph")
+    assert tables == {}
+    descriptor = _build_descriptor(8, 2, 0)
+    assert fitting_routes.fitting_routes(_route_fitting(descriptor, 96)) == []
 
 
 @_GPU
-def test_spin_bond_family_couples_the_kernel_to_the_edge_direction(
+def test_a_measured_table_names_only_its_departures_from_the_rules() -> None:
+    """A measured table leaves the rules every run size that they route as measured.
+
+    The measurement runs on the bands of the device with a stand-in for the
+    timing, under which the per-layer kernels win from a run size inside one
+    band of throughput runs to the end of the next band, and the route of the
+    rules wins at every other run size, the per-layer kernels of the smallest
+    runs among them. The table holds the rules wherever that winner is the
+    route that graph_fitting_route_rule gives, and the per-layer kernels over
+    the departure, whose start inside its band the measurement bisects for.
+    The rules and the stand-in change their route only at the probed sizes,
+    the bounds of the bands and of the departure.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+    from deepmd.pt_expt.kernels.graph_fitting import (
+        fitting_operator_arguments,
+    )
+
+    fitting = _route_fitting(_build_descriptor(8, 2, 0), 192)
+    ws = list(fitting_operator_arguments(fitting).weights)
+    device = torch.device("cuda", torch.cuda.current_device())
+    largest, block = fitting_routes.MAX_RUN_NODES, fitting_routes.BLOCK
+    bands = [
+        int(bound)
+        for bound in torch.ops.deepmd.graph_fitting_route_bands(
+            device.index, int(ws[0].shape[1]), int(ws[0].shape[0]), largest
+        )
+    ]
+    multiprocessors = torch.cuda.get_device_properties(device).multi_processor_count
+    band = next(
+        index
+        for index in range(1, len(bands) - 1)
+        if bands[index - 1] >= 32 * multiprocessors
+        and bands[index] - bands[index - 1] >= 4 * block
+    )
+    start = block * (math.ceil((bands[band - 1] + 1) / block) + 1)
+    end = bands[band + 1]
+
+    def winner(node_count: int, rule: int, choices: list[int]) -> int:
+        return fitting_routes.PER_LAYER if start < node_count <= end else rule
+
+    bounds, choices = fitting_routes._measure(
+        types.SimpleNamespace(device=device, winner=winner), ws, largest
+    )
+    probes = {1, *(size + d for size in (*bands, start, end) for d in (0, 1))}
+    for node_count in sorted(size for size in probes if size <= largest):
+        rule = int(torch.ops.deepmd.graph_fitting_route_rule(ws, node_count))
+        fastest = winner(node_count, rule, [])
+        recorded = next(
+            choice
+            for bound, choice in zip([*bounds, largest], choices, strict=True)
+            if node_count <= bound
+        )
+        expected = fitting_routes.RULES if fastest == rule else fastest
+        assert recorded == expected, node_count
+    assert fitting_routes.PER_LAYER in choices
+
+
+def _route_stand_in(
+    alternating: dict[int, float | tuple[float, ...]],
+    trains: dict[str, dict[int, tuple]],
+    type_count: int = 3,
+) -> tuple[Any, list[int], list[tuple[str, int]]]:
+    """Return a route bench whose timings are those of the given stand-ins.
+
+    ``alternating`` holds the replay time of every route in the alternating
+    replays of a sample, one time for every capture or one per capture after
+    the untimed round. ``trains`` holds the times of the trains of one route
+    at a time, one per capture, for each composition: every node of one type
+    (``"one type"``) and the node types of the sample (``"sample"``), drawn
+    among ``type_count`` types. The two lists returned with the bench record
+    the routes of the alternating replays, the untimed round included, and the
+    compositions and routes of the trains, in the order in which the bench
+    times them.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+
+    replayed: list[int] = []
+    trained: list[tuple[str, int]] = []
+    calls = 0
+
+    def replay(graphs: dict, replays: int, times: dict[int, list[float]]) -> None:
+        nonlocal calls
+        # The first call of a sample is its untimed round, which takes the
+        # times of the first capture; call k > 0 times capture k - 1.
+        capture = max(calls - 1, 0)
+        calls += 1
+        for choice in graphs:
+            replayed.append(choice)
+            value = alternating[choice]
+            ms = value[capture] if isinstance(value, tuple) else value
+            times[choice].extend([ms] * replays)
+
+    def capture(inputs: tuple, choices: list[int], rotation: int) -> dict:
+        nodes = "sample" if bool(inputs[1].any()) else "one type"
+        return {choice: (nodes, choice) for choice in choices}
+
+    def train_ms(graph: tuple[str, int]) -> float:
+        trained.append(graph)
+        nodes, choice = graph
+        return remaining[nodes][choice].pop(0)
+
+    remaining = {
+        nodes: {choice: list(times) for choice, times in routes.items()}
+        for nodes, routes in trains.items()
+    }
+    bench = object.__new__(fitting_routes._Bench)
+    bench.samples = bench.checked = bench.overturned = bench.replays = 0
+    bench.types = type_count
+    bench.one_type = torch.zeros(64, dtype=torch.long)
+    bench._inputs = lambda node_count: (None, torch.arange(64) % type_count, None)
+    bench._capture = capture
+    bench._replay_ms = lambda graph: 5.0
+    bench._replay = replay
+    bench._train_ms = train_ms
+    return bench, replayed, trained
+
+
+def _alternating_ms(*gains: float) -> tuple[float, ...]:
+    """Return per-capture replay times whose gains over 1 ms are ``gains``, in percent."""
+    return tuple(1.0 / (1.0 + gain / 100.0) for gain in gains)
+
+
+def _train_order(nodes: str, rule: int, route: int) -> list[tuple[str, int]]:
+    """Return the trains of the two captures of one composition, in their order."""
+    return [(nodes, choice) for choice in (rule, route, route, rule)]
+
+
+@pytest.mark.parametrize(
+    ("alternating_ms", "one_type_ms", "sample_ms", "per_layer_wins", "compositions"),
+    [
+        (1.05, (0.95, 0.95), (0.95, 0.95), False, 0),
+        (0.95, (1.10, 1.10), (0.95, 0.95), False, 1),
+        (0.95, (1.00, 1.04), (0.95, 0.95), False, 1),
+        (0.97, (0.95, 0.95), (1.02, 1.02), False, 2),
+        (0.97, (1.005, 1.005), (1.008, 1.008), True, 2),
+        (0.97, (0.98, 1.03), (1.00, 1.00), True, 2),
+    ],
+    ids=[
+        "slower-alternating",
+        "slower-one-type",
+        "median-beyond-the-tie",
+        "slower-sample",
+        "within-the-tie",
+        "median-within-the-tie",
+    ],
+)
+def test_a_departure_holds_on_trains_of_both_node_compositions(
+    alternating_ms: float,
+    one_type_ms: tuple,
+    sample_ms: tuple,
+    per_layer_wins: bool,
+    compositions: int,
+) -> None:
+    """A route that wins the alternating replays departs where it also holds on trains.
+
+    The stand-in times the route of the rules at 1 ms per replay and the
+    per-layer kernels at ``alternating_ms`` in the alternating replays of a
+    sample. On the trains of one route at a time, two captures for each
+    composition, it times the route of the rules at 1 ms and the per-layer
+    kernels at ``one_type_ms`` with every node of one type and at
+    ``sample_ms`` with the node types of the sample. A board at its power limit
+    can make a route that wins the alternating replays lose on trains, by an
+    amount that depends on the composition, because it sets the clock of a
+    route from the power that the route draws, which the node types of the
+    system change.
+
+    A route that clearly loses the alternating replays is settled after one
+    capture and never timed on trains, and a route that wins them is timed on
+    three captures at least. It departs from the rules where its median gain
+    on the trains of each composition is at least -1%, the tie, and the check
+    stops at the first composition on which it loses by more. The two captures
+    of a composition time the two routes in opposite orders.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+
+    rule, per_layer = 7, fitting_routes.PER_LAYER
+    bench, replayed, trained = _route_stand_in(
+        {per_layer: alternating_ms, rule: 1.0},
+        {
+            "one type": {per_layer: one_type_ms, rule: (1.0, 1.0)},
+            "sample": {per_layer: sample_ms, rule: (1.0, 1.0)},
+        },
+    )
+    expected = per_layer if per_layer_wins else rule
+    assert bench.winner(4096, rule, [per_layer, rule]) == expected
+    assert bench.samples == 1
+    # The untimed round and the captures of the sample, two routes each.
+    captures = 1 if alternating_ms > 1.0 else 3
+    assert len(replayed) == 2 * (1 + captures)
+    checked = int(compositions > 0)
+    assert (bench.checked, bench.overturned) == (checked, checked - per_layer_wins)
+    order = _train_order("one type", rule, per_layer) + _train_order(
+        "sample", rule, per_layer
+    )
+    assert trained == order[: 4 * compositions]
+
+
+@pytest.mark.parametrize(
+    ("per_layer_gains", "shape_gains", "expected"),
+    [
+        ((0.5,) * 7, None, "rules"),
+        ((2.0, 2.0, 10.0, 2.0, 2.0, 10.0, 2.0), None, "rules"),
+        ((6.0, 5.0, 0.0, 6.0, 5.0, 0.0, 5.0), None, "per-layer"),
+        ((2.5, 2.5, -6.0, 2.5, 2.5, -6.0, 2.5), (2.0,) * 7, "shape"),
+    ],
+    ids=[
+        "within-the-tie",
+        "within-the-half-range",
+        "beyond-the-half-range",
+        "runner-up-of-a-wide-best",
+    ],
+)
+def test_a_candidate_exceeds_its_own_bound(
+    per_layer_gains: tuple, shape_gains: tuple | None, expected: str
+) -> None:
+    """A route is a candidate where its median gain exceeds its own bound.
+
+    The stand-in times the route of the rules at 1 ms per replay in the
+    alternating replays of a sample and the per-layer kernels, and in the
+    last case a fused shape, at the given per-capture gains over it, in
+    percent; the trains time every route at 1 ms. The bound of a route is the
+    tie of 1%, or the half-range of its own per-capture gains where that is
+    larger. A steady gain of 0.5% within the tie and a median gain of 2% within
+    a half-range of 4% leave the run size to the rules, a median gain of 5%
+    beyond a half-range of 3% departs, and a fused shape of a steady 2% departs
+    where the per-layer kernels do not, although their median gain of 2.5% is
+    larger, since the half-range of their gains is 4.25%.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+
+    rule, shape, per_layer = 7, 5, fitting_routes.PER_LAYER
+    alternating = {per_layer: _alternating_ms(*per_layer_gains), rule: 1.0}
+    if shape_gains is not None:
+        alternating[shape] = _alternating_ms(*shape_gains)
+    level = dict.fromkeys((per_layer, shape, rule), (1.0,) * 4)
+    bench, _, _ = _route_stand_in(alternating, {"one type": level, "sample": level})
+    routes = {"rules": rule, "per-layer": per_layer, "shape": shape}
+    assert bench.winner(4096, rule, [*alternating]) == routes[expected]
+    assert bench.checked == int(expected != "rules")
+
+
+@pytest.mark.parametrize(
+    ("per_layer_gains", "shape_gains", "expected", "stages"),
+    [
+        (
+            (2.2, 2.2, -6.0, 0.0, 0.0, 0.0, 0.0),
+            (1.9, 2.0, 2.1, 0.0, 0.0, 0.0, 0.0),
+            "rules",
+            3,
+        ),
+        ((6.0,) * 7, (-6.0,) * 7, "per-layer", 2),
+        ((-6.0,) * 7, (6.0,) * 7, "shape", 2),
+    ],
+    ids=[
+        "runner-up-near-its-bound",
+        "departure-of-the-first-route",
+        "departure-of-the-second-route",
+    ],
+)
+def test_a_sample_settles_every_rival_before_it_stops(
+    per_layer_gains: tuple, shape_gains: tuple, expected: str, stages: int
+) -> None:
+    """A sample takes captures until every route besides the rules' is settled.
+
+    The stand-in times the route of the rules at 1 ms per replay and the
+    per-layer kernels and a fused shape, given to the bench in that order, at
+    the given per-capture gains over it, in percent; the trains time every
+    route at 1 ms. In the first case, after three captures the per-layer
+    kernels, whose median gain is the largest, lie clearly below their bound,
+    a half-range of 4.1%, while the fused shape lies 1% above its bound,
+    within the margin of that stage: the sample takes the captures of the
+    third stage, over which both routes tie with the rules, and leaves the run
+    size to them. In the other two cases one route lies clearly below its
+    bound after one capture and the other clearly above its own, the first
+    route given to the bench in one case and the second in the other: the
+    sample still takes the captures of the second stage before the route
+    above its bound departs.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+
+    rule, shape, per_layer = 7, 5, fitting_routes.PER_LAYER
+    level = dict.fromkeys((per_layer, shape, rule), (1.0,) * 4)
+    bench, replayed, _ = _route_stand_in(
+        {
+            per_layer: _alternating_ms(*per_layer_gains),
+            shape: _alternating_ms(*shape_gains),
+            rule: 1.0,
+        },
+        {"one type": level, "sample": level},
+    )
+    routes = {"rules": rule, "per-layer": per_layer, "shape": shape}
+    assert bench.winner(4096, rule, [per_layer, shape, rule]) == routes[expected]
+    # The untimed round and the captures of the stages, three routes each.
+    captures = sum(fitting_routes.STAGE_CAPTURES[:stages])
+    assert len(replayed) == 3 * (1 + captures)
+
+
+@pytest.mark.parametrize(
+    ("one_type_ms", "expected"),
+    [((1.03, 0.97), 1), ((1.03, 1.02), None), ((1.03, 1.03, 0.99), 2)],
+    ids=["second-holds", "none-holds", "third-holds"],
+)
+def test_the_candidates_are_tried_in_the_order_of_their_gains(
+    one_type_ms: tuple[float, ...], expected: int | None
+) -> None:
+    """The candidates of a sample are checked on trains in the order of their gains.
+
+    The stand-in times the per-layer kernels and one or two fused shapes at
+    0.95, 0.96 and 0.97 ms in the alternating replays of a sample, against
+    1 ms for the route of the rules, so that every one exceeds the tie and
+    their gains fall in that order; the bench receives the routes in the
+    reverse order. On the trains with every node of one type candidate ``i``
+    takes ``one_type_ms[i]`` against 1 ms for the route of the rules, so that
+    the per-layer kernels lose 3% there, and with the node types of the sample
+    every candidate takes 0.97 ms. The first candidate that holds on both
+    compositions departs from the rules, and the rules keep the run size where
+    none holds; the sample counts as overturned either way.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+
+    rule = 7
+    candidates = [fitting_routes.PER_LAYER, 5, 6][: len(one_type_ms)]
+    alternating = {choice: 0.95 + 0.01 * i for i, choice in enumerate(candidates)}
+    one_type = dict(zip(candidates, ((ms, ms) for ms in one_type_ms), strict=True))
+    bench, _, trained = _route_stand_in(
+        {**alternating, rule: 1.0},
+        {
+            "one type": {**one_type, rule: (1.0,) * (2 * len(candidates))},
+            "sample": {**dict.fromkeys(candidates, (0.97, 0.97)), rule: (1.0, 1.0)},
+        },
+    )
+    winner = rule if expected is None else candidates[expected]
+    assert bench.winner(4096, rule, [*reversed(candidates), rule]) == winner
+    assert (bench.checked, bench.overturned) == (1, 1)
+    tried = candidates if expected is None else candidates[: expected + 1]
+    order = [train for c in tried for train in _train_order("one type", rule, c)]
+    if expected is not None:
+        order += _train_order("sample", rule, winner)
+    assert trained == order
+
+
+@pytest.mark.parametrize(
+    ("rule_ms", "rule_train_ms", "per_layer_wins", "compositions"),
+    [(1.125, 1.0, False, 0), (1.25, 0.875, True, 2)],
+    ids=["gain-at-its-bound", "train-gain-at-the-tie"],
+)
+def test_the_bound_and_the_tie_are_strict_and_inclusive(
     monkeypatch: pytest.MonkeyPatch,
+    rule_ms: float,
+    rule_train_ms: float,
+    per_layer_wins: bool,
+    compositions: int,
 ) -> None:
-    """The compiled spin path reads the edge direction, and only through it.
+    """A gain equal to its bound is no candidate, and a train loss equal to the tie holds.
+
+    With a tie of 1/8, which the times of the stand-in represent exactly, the
+    per-layer kernels take 1 ms in the alternating replays and on the trains.
+    A route of the rules at 1.125 ms gives them a steady gain of exactly the
+    tie, their bound, which a candidate has to exceed, so the run size stays
+    with the rules. A route of the rules at 1.25 ms in the alternating replays
+    and at 0.875 ms on the trains gives them a candidacy and a median gain of
+    exactly minus the tie on the trains of both compositions, at which a
+    candidate holds.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
 
-    Rotating the neighbourhood while holding the moments fixed leaves every
-    spin family except the bond-projected one invariant, so a kernel that
-    omitted that family would return an unchanged readout. Rotating the
-    moments alongside the geometry restores the invariance, which is what
-    separates a genuine bond coupling from a defect in the geometric block.
+    monkeypatch.setattr(fitting_routes, "TIE", 0.125)
+    rule, per_layer = 7, fitting_routes.PER_LAYER
+    level = {per_layer: (1.0, 1.0), rule: (rule_train_ms, rule_train_ms)}
+    bench, _, trained = _route_stand_in(
+        {per_layer: 1.0, rule: rule_ms}, {"one type": level, "sample": level}
+    )
+    expected = per_layer if per_layer_wins else rule
+    assert bench.winner(4096, rule, [per_layer, rule]) == expected
+    assert len(trained) == 4 * compositions
+
+
+def test_a_margin_equal_to_that_of_its_stage_takes_the_next_stage(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A sample decides only where every margin lies strictly beyond that of its stage.
+
+    With a tie of 1/8 and stage margins of 1/4 and 1/8, which the times of the
+    stand-in represent exactly, the route of the rules takes 1.25 ms and the
+    per-layer kernels 1 ms per replay, a steady gain of 1/4. Its margin, the
+    gain less its bound, is 1/8 after every stage. The first stage does not
+    decide, since the route exceeds its bound, and the second does not either,
+    since the margin equals that of the stage, so the sample takes the
+    captures of the last stage before the per-layer kernels depart.
     """
-    descriptor = _build_spin_descriptor(16)
-    graph, atype = _build_graph(descriptor, canonical=True)
-    generator = torch.Generator(device="cuda").manual_seed(11)
-    spin = torch.randn(
-        atype.shape[0],
-        3,
-        dtype=torch.float32,
-        device="cuda",
-        generator=generator,
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
     )
-    angle = 0.4
-    rotation = torch.tensor(
-        [
-            [float(np.cos(angle)), -float(np.sin(angle)), 0.0],
-            [float(np.sin(angle)), float(np.cos(angle)), 0.0],
-            [0.0, 0.0, 1.0],
-        ],
-        dtype=torch.float32,
-        device="cuda",
+
+    monkeypatch.setattr(fitting_routes, "TIE", 0.125)
+    monkeypatch.setattr(fitting_routes, "STAGE_MARGINS", (0.25, 0.125))
+    rule, per_layer = 7, fitting_routes.PER_LAYER
+    level = {per_layer: (1.0, 1.0), rule: (1.0, 1.0)}
+    bench, replayed, _ = _route_stand_in(
+        {per_layer: 1.0, rule: 1.25}, {"one type": level, "sample": level}
     )
+    assert bench.winner(4096, rule, [per_layer, rule]) == per_layer
+    # The untimed round and the captures of the three stages, two routes each.
+    assert len(replayed) == 2 * (1 + sum(fitting_routes.STAGE_CAPTURES))
 
-    def run(edge_vec: torch.Tensor, moment: torch.Tensor) -> torch.Tensor:
-        output, _ = descriptor.call_graph(
-            dataclasses.replace(graph, edge_vec=edge_vec),
-            atype,
-            spin=moment,
-        )
-        return output
 
-    portable = run(graph.edge_vec, spin)
-    descriptor.enable_compression(0.5)
-    monkeypatch.setenv("DP_CUDA_INFER", "1")
-    upright = run(graph.edge_vec, spin)
-    _assert_dispatched(upright, portable)
-    rotated = run(graph.edge_vec @ rotation.T, spin)
-    assert not torch.allclose(rotated, upright, atol=1e-4)
-    covariant = run(graph.edge_vec @ rotation.T, spin @ rotation.T)
-    torch.testing.assert_close(covariant, upright, atol=3e-5, rtol=3e-5)
+def test_the_candidates_are_ordered_by_their_gain_rather_than_their_margin() -> None:
+    """Of two candidates, the one of the larger median gain is checked first.
+
+    The stand-in times the route of the rules at 1 ms per replay, the
+    per-layer kernels at per-capture gains of 6%, 5% and 0%, a median gain of
+    5% above a bound of 3%, and a fused shape at a steady gain of 4%, above a
+    bound of 1%, so that the fused shape lies farther above its bound while
+    the per-layer kernels gain more. The trains time every route at 1 ms, so
+    both candidates hold there, and the per-layer kernels, checked first, take
+    the run size.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+
+    rule, shape, per_layer = 7, 5, fitting_routes.PER_LAYER
+    level = dict.fromkeys((per_layer, shape, rule), (1.0,) * 4)
+    bench, _, trained = _route_stand_in(
+        {
+            per_layer: _alternating_ms(6.0, 5.0, 0.0, 6.0, 5.0, 0.0, 5.0),
+            shape: _alternating_ms(*(4.0,) * 7),
+            rule: 1.0,
+        },
+        {"one type": level, "sample": level},
+    )
+    assert bench.winner(4096, rule, [shape, per_layer, rule]) == per_layer
+    assert trained == _train_order("one type", rule, per_layer) + _train_order(
+        "sample", rule, per_layer
+    )
 
 
 @_GPU
-@pytest.mark.parametrize("bridged", [False, True])
-@pytest.mark.parametrize("spin_conditioned", [False, True])
-def test_backward_operator_satisfies_its_schema(
-    spin_conditioned: bool,
-    bridged: bool,
+@pytest.mark.parametrize(
+    ("train_span_ms", "floor"), [(10.0, False), (0.5, True)], ids=["span", "floor"]
+)
+def test_a_train_follows_its_lead_in_and_returns_the_time_of_one_replay(
+    monkeypatch: pytest.MonkeyPatch, train_span_ms: float, floor: bool
 ) -> None:
-    """The backward operator declares four independent results.
+    """A train replays its route for the lead-in, then times a span of replays.
+
+    Captured CUDA graphs of sleep kernels stand in for the operator: one
+    replays during the lead-in, whose replays the bench times through
+    ``_replay_ms``, and one of four times its length during the train, so that
+    the time that the train returns tells its own replays from those of the
+    lead-in also where other processes take time slices of a shared device and
+    so lengthen the lead-in. With a lead-in of 20 ms of device time in chunks of
+    5 ms, the first replay sets the size of the chunks, the lead-in replays the
+    route in chunks until their device time reaches 20 ms, and the train then
+    takes the replays that span ``train_span_ms`` at the mean replay time of
+    the lead-in, nine at least: a span of 10 ms takes more than nine replays, a
+    span of 0.5 ms nine. The bench counts every replay, and the train returns
+    the median time of its own replays, which the test compares with the
+    median of single replays of the same graph; a median, unlike a mean, stays
+    in place under the time slices of other processes.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
 
-    All four are unannotated, so any two of them sharing storage would be an
-    alias the schema does not describe, which is undefined under
-    functionalization. ``opcheck`` decides that mechanically, including on the
-    spin-free path without a pair potential, where three of the four are
-    absent and an allocation shared between them would otherwise go unnoticed.
+    monkeypatch.setattr(fitting_routes, "LEAD_IN_MS", 20.0)
+    monkeypatch.setattr(fitting_routes, "LEAD_IN_CHUNK_MS", 5.0)
+    monkeypatch.setattr(fitting_routes, "TRAIN_MS", train_span_ms)
+    device = torch.device("cuda", torch.cuda.current_device())
+
+    class Counted:
+        """Replays the lead-in graph or the train graph and counts the replays."""
+
+        def __init__(
+            self, lead_in: torch.cuda.CUDAGraph, train: torch.cuda.CUDAGraph
+        ) -> None:
+            self.lead_in, self.train = lead_in, train
+            self.leading, self.count = False, 0
+
+        def replay(self) -> None:
+            self.count += 1
+            (self.lead_in if self.leading else self.train).replay()
+
+    bench = object.__new__(fitting_routes._Bench)
+    bench.device, bench.replays = device, 0
+    timed: list[tuple[int, float]] = []
+    replay_ms = bench._replay_ms
+
+    def lead_in_replay_ms(graph: Counted, count: int = 1) -> float:
+        graph.leading = True
+        timed.append((count, replay_ms(graph, count)))
+        graph.leading = False
+        return timed[-1][1]
+
+    bench._replay_ms = lead_in_replay_ms
+    with torch.cuda.device(device):
+        stream = torch.cuda.Stream(device)
+        with torch.cuda.stream(stream):
+            torch.cuda._sleep(1000)
+        torch.cuda.synchronize(device)
+        graphs = []
+        for cycles in (250_000, 1_000_000):
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                torch.cuda._sleep(cycles)
+            graphs.append(graph)
+        counted = Counted(*graphs)
+        train_ms = bench._train_ms(counted)
+        single_ms = statistics.median(replay_ms(graphs[1]) for _ in range(15))
+    counts, times = zip(*timed, strict=True)
+    chunk = max(1, math.ceil(fitting_routes.LEAD_IN_CHUNK_MS / times[0]))
+    assert counts == (1, *(chunk,) * (len(counts) - 1))
+    assert sum(times[:-1]) < fitting_routes.LEAD_IN_MS <= sum(times)
+    span = math.ceil(fitting_routes.TRAIN_MS * sum(counts) / sum(times))
+    if floor:
+        assert span < fitting_routes.TRAIN_REPLAYS
+    else:
+        assert span > fitting_routes.TRAIN_REPLAYS
+    assert counted.count - sum(counts) == max(fitting_routes.TRAIN_REPLAYS, span)
+    assert bench.replays == counted.count
+    assert 0.8 * single_ms <= train_ms <= 1.25 * single_ms
+
+
+@pytest.mark.parametrize("type_count", [1, 2], ids=["one-type", "two-types"])
+def test_the_trains_time_each_distinct_composition_once(type_count: int) -> None:
+    """The trains time each distinct composition of a sample once.
+
+    The node types of the sample of a model of one node type are those of the
+    composition with every node of one type, so the bench times that
+    composition alone; the sample of a model of two node types holds both
+    types, and the bench times both compositions.
     """
-    from deepmd.dpmodel.atomic_model.inner_potential import (
-        InnerPotential,
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
     )
 
-    descriptor = _build_spin_descriptor(8) if spin_conditioned else _build_descriptor(8)
-    graph, atype = _build_graph(descriptor, canonical=True)
-    arguments = _arguments(descriptor, graph, atype)
-    pair = (
-        (
-            torch.as_tensor(InnerPotential(["O", "H"]).pair_table, device="cuda"),
-            torch.ones(atype.shape[0], dtype=torch.float64, device="cuda"),
-        )
-        if bridged
-        else (graph.edge_vec.new_empty(0),) * 2
+    rule, per_layer = 7, fitting_routes.PER_LAYER
+    level = {per_layer: (0.95, 0.95), rule: (1.0, 1.0)}
+    bench, _, trained = _route_stand_in(
+        {per_layer: 0.95, rule: 1.0},
+        {"one type": level, "sample": level},
+        type_count=type_count,
     )
-    if spin_conditioned:
-        arguments = _with_spin(
-            arguments,
-            torch.randn(atype.shape[0], 3, dtype=torch.float32, device="cuda"),
-        )
-    output, state = torch.ops.deepmd.dpa4c_graph_compress(
-        graph.edge_vec,
-        *arguments,
+    assert bench.winner(4096, rule, [per_layer, rule]) == per_layer
+    assert (bench.checked, bench.overturned) == (1, 0)
+    order = _train_order("one type", rule, per_layer)
+    if type_count > 1:
+        order += _train_order("sample", rule, per_layer)
+    assert trained == order
+
+
+@pytest.mark.parametrize(
+    ("one_type_ms", "sample_ms", "second_wins"),
+    [
+        ((1.10, 1.10), (1.0, 1.0), True),
+        ((1.012, 0.995), (1.0, 1.0), True),
+        ((0.995, 0.995), (1.03, 1.03), True),
+        ((0.97, 0.97), (1.0, 1.0), False),
+        ((1.02, 1.02), (0.95, 0.95), False),
+        ((1.0, 1.0), (1.0, 1.0), False),
+    ],
+    ids=["clear", "within-range", "mean", "first-faster", "mean-against", "tie"],
+)
+def test_two_routes_besides_the_rules_are_decided_on_trains(
+    one_type_ms: tuple, sample_ms: tuple, second_wins: bool
+) -> None:
+    """Two routes that are not the rules' are decided on trains alone.
+
+    The bisection of a band compares the routes that won at its two ends,
+    neither of them that of the rules, and the bench runs no alternating
+    replays for such a pair. The stand-in times the per-layer kernels, the
+    first route, faster than a fused shape in the alternating replays, and on
+    the trains at ``one_type_ms`` with every node of one type and at
+    ``sample_ms`` with the node types of the sample, against 1 ms for the
+    fused shape. The fused shape takes the run size where the mean over the two
+    compositions of its median gains over the per-layer kernels is positive,
+    and the per-layer kernels keep it otherwise, a tie included.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
     )
-    torch.library.opcheck(
-        torch.ops.deepmd.dpa4c_graph_compress_backward.default,
-        (torch.ones_like(output), state, graph.edge_vec, *arguments, *pair),
+
+    rule, shape, per_layer = 7, 5, fitting_routes.PER_LAYER
+    bench, replayed, trained = _route_stand_in(
+        {per_layer: 0.95, shape: 1.0},
+        {
+            "one type": {per_layer: one_type_ms, shape: (1.0, 1.0)},
+            "sample": {per_layer: sample_ms, shape: (1.0, 1.0)},
+        },
+    )
+    expected = shape if second_wins else per_layer
+    assert bench.winner(4096, rule, [per_layer, shape]) == expected
+    assert (bench.samples, bench.checked, bench.overturned) == (1, 1, 0)
+    assert replayed == []
+    assert trained == _train_order("one type", per_layer, shape) + _train_order(
+        "sample", per_layer, shape
     )
 
 
 @_GPU
-def test_registered_autograd_refuses_the_magnetic_moment() -> None:
-    """A direct operator call cannot silently lose the magnetic force.
+def test_the_route_tuner_measures_a_table_that_changes_no_output(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The tuner holds a table of the device under which no output changes.
 
-    The operator emits that cotangent in two pieces and the per-edge piece is
-    reduced onto source nodes through the source CSR, which the schema does
-    not carry. The registration therefore cannot close the magnetic force, and
-    refusing is the only alternative to reporting a vanishing one.
+    The routing model is measured over the run sizes of a node tile of 16,384.
+    The table carries the fingerprint of the device and increasing bounds
+    below the tile, and names on every interval the rules or a route that the
+    operator takes there. On a bulk system at the upper end of every interval,
+    the operator gives the outputs of the rules bit for bit under it.
     """
-    descriptor = _build_spin_descriptor(8)
-    graph, atype = _build_graph(descriptor, canonical=True)
-    spin = torch.randn(
-        atype.shape[0],
-        3,
-        dtype=torch.float32,
-        device="cuda",
-        requires_grad=True,
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
     )
-    output, _state = torch.ops.deepmd.dpa4c_graph_compress(
-        graph.edge_vec,
-        *_with_spin(_arguments(descriptor, graph, atype), spin),
+    from deepmd.pt_expt.kernels.dpa4c.canonical import (
+        dpa4c_canonical_compress_energy_force,
     )
-    with pytest.raises(RuntimeError, match="cannot differentiate its magnetic"):
-        output.sum().backward()
+    from deepmd.pt_expt.kernels.graph_fitting import (
+        fitting_operator_arguments,
+    )
+
+    tables: dict = {}
+    monkeypatch.setattr(fitting_routes, "_ROUTE_TABLES", tables)
+    monkeypatch.setenv("DP_NODE_TILE", "16384")
+    device = torch.device("cuda", torch.cuda.current_device())
+    model = _route_model()
+    fitting_routes.tune_fitting_routes(model, 0, device)
+    (table,) = tables.values()
+    table = list(table)
+    fingerprint = list(fitting_routes._fingerprint(device))
+    header = 1 + len(fingerprint)
+    assert table[:header] == [fitting_routes.ROUTE_TABLE_MAGIC, *fingerprint]
+    intervals = table[header]
+    bounds = table[header + 1 : header + intervals]
+    choices = table[header + intervals :]
+    assert len(choices) == intervals
+    assert bounds == sorted(set(bounds)) and all(0 < b < 16384 for b in bounds)
+    learned = model.atomic_model.fused_decomposition()[0]
+    descriptor = learned.descriptor.to(device).eval()
+    fitting = learned.fitting_net.to(device).eval()
+    weights = list(fitting_operator_arguments(fitting).weights)
+    bias = fitting.bias_atom_e[:, 0]
+    for node_count, choice in zip([*bounds, 16384], choices, strict=True):
+        if choice != fitting_routes.RULES:
+            resolved = torch.ops.deepmd.graph_fitting_route_resolve(
+                weights, node_count, table
+            )
+            assert resolved == choice, node_count
+        graph = fitting_routes._bulk_graph(4.0, node_count, device)
+        atype = torch.arange(node_count, device=device) % 2
+        rules = dpa4c_canonical_compress_energy_force(
+            descriptor, fitting, graph, atype, bias, fitting_routes=[]
+        )
+        routed = dpa4c_canonical_compress_energy_force(
+            descriptor, fitting, graph, atype, bias, fitting_routes=table
+        )
+        for expected, actual in zip(rules, routed, strict=True):
+            assert torch.equal(expected, actual), node_count
diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py
index 57fdd021ce..097feb5806 100644
--- a/source/tests/pt_expt/infer/test_deep_eval.py
+++ b/source/tests/pt_expt/infer/test_deep_eval.py
@@ -10,6 +10,9 @@
 import tempfile
 import unittest
 import zipfile
+from unittest import (
+    mock,
+)
 
 import numpy as np
 import torch
@@ -229,6 +232,23 @@ def test_serialize_returns_model_tree(self) -> None:
         # single bool and fail with an ambiguous truth-value error.
         np.testing.assert_equal(data, serialize_from_file(self.tmpfile.name)["model"])
 
+    def test_pte_load_is_independent_of_default_device(self) -> None:
+        """Archive loading and inference retain the recorded tensor devices."""
+        rng = np.random.default_rng(GLOBAL_SEED)
+        coords = rng.random((1, 5, 3)) * 8.0
+        cells = np.eye(3).reshape(1, 9) * 10.0
+        atom_types = np.array([0, 1, 0, 1, 0], dtype=np.int32)
+        expected = self.dp.eval(coords, cells, atom_types, atomic=True)
+        for default_device in ("meta", "cuda:9999999"):
+            with self.subTest(default_device=default_device):
+                with torch.device(default_device):
+                    dp = DeepPot(self.tmpfile.name)
+                    data = serialize_from_file(self.tmpfile.name)
+                    actual = dp.eval(coords, cells, atom_types, atomic=True)
+                np.testing.assert_equal(data["model"], self.model_data["model"])
+                for value, reference in zip(actual, expected, strict=True):
+                    np.testing.assert_allclose(value, reference, rtol=1e-10, atol=1e-10)
+
     def test_eval_consistency(self) -> None:
         """Test that DeepPot.eval gives same results as direct model forward."""
         rng = np.random.default_rng(GLOBAL_SEED)
@@ -2342,22 +2362,53 @@ def _vesin_available(self) -> bool:
         return is_vesin_torch_available()
 
     def test_default_is_auto(self) -> None:
-        dp = DeepPot(self.tmpfile.name)
-        self.assertEqual(dp.deep_eval._use_vesin, self._vesin_available())
+        from deepmd.pt_expt.utils import (
+            auto_neighbor_list,
+        )
+
+        for vesin_available, nv_available, expected in (
+            (False, False, False),
+            (True, False, True),
+            (False, True, True),
+            (True, True, True),
+        ):
+            with (
+                self.subTest(vesin=vesin_available, nv=nv_available),
+                mock.patch.object(
+                    auto_neighbor_list,
+                    "is_vesin_torch_available",
+                    return_value=vesin_available,
+                ),
+                mock.patch.object(
+                    auto_neighbor_list, "is_nv_available", return_value=nv_available
+                ),
+            ):
+                dp = DeepPot(self.tmpfile.name)
+                self.assertEqual(dp.deep_eval._nlist_builder is not None, expected)
 
     def test_native_disables_vesin(self) -> None:
         dp = DeepPot(self.tmpfile.name, nlist_backend="native")
-        self.assertFalse(dp.deep_eval._use_vesin)
+        self.assertIsNone(dp.deep_eval._nlist_builder)
 
     def test_invalid_raises(self) -> None:
         with self.assertRaises(ValueError):
             DeepPot(self.tmpfile.name, nlist_backend="bogus")
 
-    def test_vesin_requested_but_unavailable_raises(self) -> None:
-        if self._vesin_available():
-            self.skipTest("vesin.torch is installed; cannot test the missing path")
-        with self.assertRaises(ImportError):
-            DeepPot(self.tmpfile.name, nlist_backend="vesin")
+    def test_requested_backend_unavailable_raises(self) -> None:
+        from deepmd.pt_expt.utils import (
+            auto_neighbor_list,
+        )
+
+        for method, probe, message in (
+            ("vesin", "is_vesin_torch_available", r"vesin\[torch\]"),
+            ("nv", "is_nv_available", "nvalchemi-toolkit-ops"),
+        ):
+            with (
+                self.subTest(method=method),
+                mock.patch.object(auto_neighbor_list, probe, return_value=False),
+                self.assertRaisesRegex(ImportError, message),
+            ):
+                DeepPot(self.tmpfile.name, nlist_backend=method)
 
     def test_vesin_matches_native(self) -> None:
         if not self._vesin_available():
@@ -2376,6 +2427,25 @@ def test_vesin_matches_native(self) -> None:
                     a, b, rtol=1e-10, atol=1e-10, err_msg=f"vesin vs native: {name}"
                 )
 
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
+    def test_nv_matches_native_multiframe(self) -> None:
+        from deepmd.pt.utils.nv_nlist import (
+            is_nv_available,
+        )
+
+        if not is_nv_available():
+            self.skipTest("nvalchemi-toolkit-ops is required")
+        rng = np.random.default_rng(GLOBAL_SEED + 21)
+        coords = rng.random((3, 6, 3)) * 8.0
+        atom_types = np.arange(6, dtype=np.int32) % self.nt
+        cells = np.tile(np.eye(3).reshape(1, 9) * 10.0, (3, 1))
+        native = DeepPot(self.tmpfile.name, nlist_backend="native")
+        accelerated = DeepPot(self.tmpfile.name, nlist_backend="nv")
+        expected = native.eval(coords, cells, atom_types, atomic=True)
+        actual = accelerated.eval(coords, cells, atom_types, atomic=True)
+        for value, reference in zip(actual, expected, strict=True):
+            np.testing.assert_allclose(value, reference, rtol=1e-9, atol=1e-9)
+
     def test_eval_descriptor_vesin_matches_native(self) -> None:
         """eval_descriptor bypasses forward_common_lower's format_nlist, so the
         vesin builder must apply the same type-distinguishing as the native
@@ -2542,6 +2612,8 @@ class TestCompressedChargeStateRoutes(unittest.TestCase):
     answer for the state the caller asked for.
     """
 
+    lower_kind = "graph"
+
     @classmethod
     def setUpClass(cls) -> None:
         import copy
@@ -2562,7 +2634,12 @@ def setUpClass(cls) -> None:
                 "default_chg_spin": [0.0, 1.0],
                 "seed": GLOBAL_SEED,
             },
-            "fitting_net": {"neuron": [8, 8], "precision": "float32", "seed": 1},
+            "fitting_net": {
+                "neuron": [8, 8],
+                "precision": "float32",
+                "resnet_dt": cls.lower_kind == "graph",
+                "seed": 1,
+            },
         }
         model = pt_expt_get_model(copy.deepcopy(config)).to("cuda")
 
@@ -2584,6 +2661,7 @@ def setUpClass(cls) -> None:
                     )
                     * 0.3
                 )
+        cls.reference = copy.deepcopy(model).eval()
         model.atomic_model.descriptor.enable_compression(0.5)
 
         cls.tmpfile = tempfile.NamedTemporaryFile(suffix=".pt2", delete=False)
@@ -2591,11 +2669,15 @@ def setUpClass(cls) -> None:
         torch.set_default_device(None)
         try:
             deserialize_to_file(
-                cls.tmpfile.name, {"model": model.serialize()}, lower_kind="auto"
+                cls.tmpfile.name,
+                {"model": model.serialize()},
+                lower_kind="auto",
+                with_comm=False,
             )
             # The nested rebuild is loaded as an archive of its own, which
             # reads the default device.
             cls.dp = DeepPot(cls.tmpfile.name)
+            assert cls.dp.deep_eval.metadata["lower_input_kind"] == cls.lower_kind
         finally:
             torch.set_default_device("cuda:9999999")
 
@@ -2655,5 +2737,62 @@ def test_a_state_is_not_left_behind(self) -> None:
         np.testing.assert_allclose(self._routes([0.0, 1.0]), first, rtol=1e-6)
 
 
+class TestCanonicalChargeStateRoutes(TestCompressedChargeStateRoutes):
+    """Canonical archives retain charge folding and both graph input paths."""
+
+    lower_kind = "dpa4c_canonical"
+
+    def test_canonical_inputs_match_eager(self) -> None:
+        from deepmd.pt.utils.nv_nlist import (
+            is_nv_available,
+        )
+        from deepmd.pt_expt.utils import (
+            canonical_graph,
+        )
+        from deepmd.pt_expt.utils.vesin_neighbor_list import (
+            is_vesin_torch_available,
+        )
+
+        state = np.array([[2.0, 3.0]])
+        reference = self.reference(
+            torch.as_tensor(self.coords, dtype=torch.float64, device=DEVICE),
+            torch.as_tensor(self.atom_types[None], dtype=torch.int64, device=DEVICE),
+            torch.as_tensor(self.cells, dtype=torch.float64, device=DEVICE),
+            charge_spin=torch.as_tensor(state, dtype=torch.float64, device=DEVICE),
+        )
+        self.assertGreater(float(reference["force"].abs().max()), 1e-8)
+        for method, available in (
+            ("dense", True),
+            ("nv", is_nv_available()),
+            ("vesin", is_vesin_torch_available()),
+        ):
+            with self.subTest(method=method):
+                if not available:
+                    self.skipTest(f"{method} is not installed")
+                dp = DeepPot(
+                    self.tmpfile.name,
+                    auto_batch_size=False,
+                    neighbor_graph_method=method,
+                )
+                with mock.patch.object(
+                    canonical_graph,
+                    "canonical_graph_from_neighbor_graph",
+                    wraps=canonical_graph.canonical_graph_from_neighbor_graph,
+                ) as convert:
+                    actual = dp.eval(
+                        self.coords, self.cells, self.atom_types, charge_spin=state
+                    )
+                self.assertEqual(convert.call_count, 0 if method == "vesin" else 1)
+                for value, key in zip(
+                    actual[:3], ("energy", "force", "virial"), strict=True
+                ):
+                    np.testing.assert_allclose(
+                        value,
+                        reference[key].detach().cpu().numpy(),
+                        rtol=1e-4,
+                        atol=1e-5,
+                    )
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py b/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py
index 00eab95d62..ef6ba81736 100644
--- a/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py
+++ b/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py
@@ -357,7 +357,7 @@ def test_nlist_controls_remain_available(self) -> None:
             output_def,
             nlist_backend="native",
         )
-        with self.assertRaisesRegex(ValueError, "only applies to graph-routed"):
+        with self.assertRaisesRegex(ValueError, "only applies to graph or edge-vector"):
             PtExptDeepEval(
                 self.pt_path,
                 output_def,
@@ -422,6 +422,93 @@ def test_unsupported_extension_raises(self) -> None:
 class TestNeighborGraphMethodResolution(unittest.TestCase):
     """Auto graph-builder selection must cover each host policy explicitly."""
 
+    def test_auto_nlist_strategy_uses_frame_count(self) -> None:
+        from deepmd.pt_expt.utils import (
+            auto_neighbor_list,
+        )
+
+        cases = (
+            ("cuda", 1, True, True, "vesin"),
+            ("cuda", 4, True, True, "nv"),
+            ("cuda", 1, True, False, "nv"),
+            ("cuda", 4, False, True, "vesin"),
+            ("cpu", 1, True, True, "vesin"),
+            ("cpu", 4, True, True, "nv"),
+            ("cuda", 1, False, False, None),
+            ("cpu", 1, False, True, "vesin"),
+            ("cpu", 1, True, False, "nv"),
+            ("cpu", 1, False, False, None),
+        )
+        for device, nf, nv_available, vesin_available, expected in cases:
+            with self.subTest(device=device, nf=nf, expected=expected):
+                coord = mock.Mock(device=torch.device(device))
+                atype = np.zeros((nf, 2), dtype=np.int64)
+                with (
+                    mock.patch.object(
+                        auto_neighbor_list, "is_nv_available", return_value=nv_available
+                    ),
+                    mock.patch.object(
+                        auto_neighbor_list,
+                        "is_vesin_torch_available",
+                        return_value=vesin_available,
+                    ),
+                ):
+                    builder = auto_neighbor_list.make_neighbor_list(
+                        "auto", coord.device, skin=1.0
+                    )
+                    if expected is None:
+                        self.assertIsNone(builder)
+                        continue
+                    with (
+                        mock.patch.object(
+                            builder._vesin, "build", return_value="vesin"
+                        ),
+                        mock.patch.object(builder._nv, "build", return_value="nv"),
+                    ):
+                        actual = builder.build(coord, atype, None, 3.0, [4])
+                self.assertEqual(actual, expected)
+
+    def test_native_edge_strategy_uses_frame_count(self) -> None:
+        """Native edge inputs use the shared carry-all strategy on every host."""
+        from deepmd.pt_expt.utils import (
+            auto_neighbor_list,
+        )
+
+        for device, nf, expected in (
+            ("cpu", 1, "vesin"),
+            ("cpu", 4, "nv"),
+            ("cuda", 1, "nv"),
+            ("cuda", 4, "nv"),
+        ):
+            with (
+                self.subTest(device=device, nf=nf),
+                mock.patch.object(
+                    auto_neighbor_list, "is_nv_available", return_value=True
+                ),
+                mock.patch.object(
+                    auto_neighbor_list, "is_vesin_torch_available", return_value=True
+                ),
+                mock.patch.object(
+                    auto_neighbor_list.VesinNeighborList, "build", return_value="vesin"
+                ) as vesin_build,
+                mock.patch.object(
+                    auto_neighbor_list.NvNeighborList, "build", return_value="nv"
+                ) as nv_build,
+            ):
+                coord = mock.Mock(device=torch.device(device))
+                atype = np.zeros((nf, 2), dtype=np.int64)
+                builder = auto_neighbor_list.make_neighbor_list(
+                    "native", coord.device, return_mode="edges"
+                )
+                actual = builder.build(
+                    coord, atype, None, 3.0, [1], return_mode="edges"
+                )
+                selected = vesin_build if expected == "vesin" else nv_build
+                selected.assert_called_once_with(
+                    coord, atype, None, 3.0, [1], return_mode="edges", pair_excl=None
+                )
+                self.assertEqual(actual, expected)
+
     def test_explicit_resolution(self) -> None:
         """Every implemented graph builder is accepted without rewriting."""
         for method in ("dense", "ase", "cell", "vesin", "nv"):
@@ -450,7 +537,8 @@ def test_auto_resolution(self) -> None:
             ("cpu", False, False, True, 1, "vesin", False),
             ("cpu", False, False, True, 4, "dense", False),
             ("cpu", False, False, False, 1, "dense", False),
-            ("cuda", True, True, True, 1, "nv", False),
+            ("cuda", True, True, True, 1, "vesin", False),
+            ("cuda", True, True, False, 1, "nv", False),
             ("cuda", True, True, True, 4, "nv", False),
             ("cuda", False, True, True, 1, "vesin", False),
             ("cuda", False, True, True, 4, "dense", True),
@@ -636,6 +724,44 @@ def test_eval_matches_public_forward(self) -> None:
                         err_msg=name,
                     )
 
+    def test_dense_introspection_survives_graph_checkpoint(self) -> None:
+        """Graph evaluation must not discard capacities used by introspection."""
+        coords = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]])
+        atom_types = np.array([0, 1], dtype=np.int32)
+        coord_t = torch.tensor(coords, dtype=torch.float64, device=DEVICE)
+        atype_t = torch.tensor(atom_types[None], dtype=torch.int64, device=DEVICE)
+        # Construct the two-atom dense list independently of the checkpoint
+        # loader: each center has one neighbor and the remaining slots pad.
+        nlist = torch.full(
+            (1, 2, sum(self.model.get_sel())), -1, dtype=torch.int64, device=DEVICE
+        )
+        nlist[0, :, 0] = torch.tensor([1, 0], device=DEVICE)
+        mapping = torch.arange(2, device=DEVICE).reshape(1, 2)
+        atomic_model = self.model.get_dp_atomic_model()
+        with torch.no_grad():
+            descriptor, rot_mat, g2, h2, _ = atomic_model.descriptor(
+                coord_t, atype_t, nlist, mapping=mapping
+            )
+            fitting = atomic_model.fitting_net
+            fitting.set_return_middle_output(True)
+            try:
+                middle = fitting(descriptor, atype_t, gr=rot_mat, g2=g2, h2=h2)[
+                    "middle_output"
+                ]
+            finally:
+                fitting.set_return_middle_output(False)
+
+        for layout, path in self.pt_paths.items():
+            with self.subTest(layout=layout):
+                dp = DeepPot(path, auto_batch_size=False)
+                for actual, expected in (
+                    (dp.eval_descriptor(coords, None, atom_types), descriptor),
+                    (dp.eval_fitting_last_layer(coords, None, atom_types), middle),
+                ):
+                    np.testing.assert_allclose(
+                        actual, expected.detach().cpu().numpy(), rtol=1e-10, atol=1e-10
+                    )
+
     def test_nlist_controls_fail_without_changing_graph_semantics(self) -> None:
         output_def = ModelOutputDef(self.model.atomic_output_def())
         path = self.pt_paths["plain"]
diff --git a/source/tests/pt_expt/kernels/test_grid_pair_train.py b/source/tests/pt_expt/kernels/test_grid_pair_train.py
index 5d8d6b9223..92f2bbe549 100644
--- a/source/tests/pt_expt/kernels/test_grid_pair_train.py
+++ b/source/tests/pt_expt/kernels/test_grid_pair_train.py
@@ -17,6 +17,16 @@
     annotations,
 )
 
+from contextlib import (
+    nullcontext,
+)
+from types import (
+    SimpleNamespace,
+)
+from typing import (
+    TYPE_CHECKING,
+)
+
 import pytest
 import torch
 from torch._dynamo.testing import (
@@ -26,6 +36,7 @@
     make_fx,
 )
 
+from deepmd.pt_expt.kernels.triton.sezm import grid_pair as grid_pair_module
 from deepmd.pt_expt.kernels.triton.sezm.grid_pair import (
     GRID_PAIR_TRITON_AVAILABLE,
     _built_in_launch_config,
@@ -39,10 +50,22 @@
     median_deviations,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+        Iterator,
+    )
+
+    from triton.compiler.compiler import (
+        CompiledKernel,
+    )
+
 pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"),
     pytest.mark.skipif(not GRID_PAIR_TRITON_AVAILABLE, reason="Triton is unavailable"),
 ]
+requires_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
 
 # ``(lmax, n_frames, n_focus, channels, n_grid)`` spanning the deployed grid
 # shapes. The slot count ``(lmax + 1)^2 * n_frames`` drives the operator's
@@ -62,6 +85,135 @@
 DRAW_SEEDS = (11, 2027, 40529)
 
 
+@pytest.fixture
+def isolated_launch_selection(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
+    """Keep selector tests independent of CUDA hardware and prior cache entries."""
+    monkeypatch.setattr(torch.cuda, "device", lambda device: nullcontext())
+    monkeypatch.setattr(torch.cuda, "get_device_name", lambda device: "test device")
+    grid_pair_module._select_launch_config.cache_clear()
+    yield
+    grid_pair_module._select_launch_config.cache_clear()
+
+
+@pytest.mark.usefixtures("isolated_launch_selection")
+@pytest.mark.parametrize("failure_stage", ["warmup", "handles"])
+def test_launch_selection_retries_and_caches(
+    failure_stage: str,
+) -> None:
+    """Resource rejection tries the next candidate and caches its admission."""
+    attempts = []
+
+    class Kernel:
+        def warmup(self, *args: object, **kwargs: object) -> SimpleNamespace:
+            attempts.append((kwargs["C_BLK"], kwargs["BLOCK_G"], kwargs["num_stages"]))
+
+            def init_handles() -> None:
+                if len(attempts) <= 2:
+                    raise grid_pair_module.triton.runtime.errors.OutOfResources(
+                        2, 1, "test"
+                    )
+
+            if failure_stage == "warmup":
+                init_handles()
+                return SimpleNamespace(_init_handles=lambda: None)
+            return SimpleNamespace(_init_handles=init_handles)
+
+    kernel = Kernel()
+    shape = ("test_kernel", 147, 64, 2, 3, False, 587, torch.float32)
+    select = grid_pair_module._select_launch_config
+    result = select(kernel, torch.device("cuda:0"), shape, (torch.float32, 2), 1)
+    assert result == (32, 32, 1)
+    assert attempts == [(32, 64, 2), (32, 32, 2), (32, 32, 1)]
+    assert (
+        select(kernel, torch.device("cuda:0"), shape, (torch.float32, 2), 1) == result
+    )
+    assert len(attempts) == 3
+
+
+@pytest.mark.usefixtures("isolated_launch_selection")
+def test_admitted_launch_preserves_runtime_arguments_and_reuses_selection(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """An admitted tile executes with each input's actual layout and pair count."""
+    probes = []
+    loaded = []
+    launches = []
+
+    class Kernel:
+        fn = SimpleNamespace(__name__="test_kernel")
+
+        def warmup(self, *args: object, **kwargs: object) -> SimpleNamespace:
+            probes.append((args, kwargs))
+            return SimpleNamespace(_init_handles=lambda: loaded.append(True))
+
+    class Launcher:
+        def __getitem__(self, grid: tuple[int, int]) -> Callable[..., None]:
+            def execute(*args: object, **kwargs: object) -> None:
+                launches.append((grid, args, kwargs))
+
+            return execute
+
+    kernel = Kernel()
+    monkeypatch.setattr(grid_pair_module, "wrap_triton", lambda candidate: Launcher())
+    for n_node in (2, 3):
+        value = torch.empty(49, n_node, 2, 192, device="cpu").permute(1, 0, 2, 3)
+        strides = value.stride()
+        grid_pair_module._launch(kernel, value, 587, 3, False, (value, *strides), 1)
+        grid, args, kwargs = launches[-1]
+        assert grid == (2 * n_node, 2)
+        assert args[0] is value
+        assert args[1:-1] == strides
+        assert args[-1] == 2 * n_node
+        assert kwargs == {
+            key: val for key, val in probes[0][1].items() if key != "grid"
+        }
+    assert len(probes) == len(loaded) == 1
+    assert len(launches) == 2
+    assert grid_pair_module._select_launch_config.cache_info().hits == 1
+
+
+@pytest.mark.usefixtures("isolated_launch_selection")
+def test_launch_rejection_is_cached_and_reported() -> None:
+    """An exhausted selector prevents execution and reuses its cached miss."""
+    attempts = []
+
+    class Kernel:
+        fn = SimpleNamespace(__name__="test_kernel")
+
+        def warmup(self, *args: object, **kwargs: object) -> None:
+            attempts.append((kwargs["C_BLK"], kwargs["BLOCK_G"], kwargs["num_stages"]))
+            raise grid_pair_module.triton.runtime.errors.OutOfResources(2, 1, "test")
+
+    kernel = Kernel()
+    select = grid_pair_module._select_launch_config
+    value = torch.empty(2, 49, 2, 192, device="cpu")
+    for _ in range(2):
+        with pytest.raises(grid_pair_module._NoViableConfig, match="P=147, C=64"):
+            grid_pair_module._launch(kernel, value, 587, 3, False, (value,), 1)
+    assert len(attempts) == len(set(attempts)) == 8
+    assert select.cache_info().hits == 1
+
+
+@pytest.mark.usefixtures("isolated_launch_selection")
+def test_launch_selection_preserves_compiler_errors() -> None:
+    """A compiler defect must neither select a fallback nor enter the cache."""
+    attempts = []
+
+    class Kernel:
+        def warmup(self, *args: object, **kwargs: object) -> None:
+            attempts.append(kwargs)
+            raise RuntimeError("compiler defect")
+
+    kernel = Kernel()
+    select = grid_pair_module._select_launch_config
+    shape = ("test_kernel", 147, 64, 2, 3, False, 587, torch.float32)
+    for _ in range(2):
+        with pytest.raises(RuntimeError, match="compiler defect"):
+            select(kernel, torch.device("cuda:0"), shape, (torch.float32, 2), 1)
+    assert len(attempts) == 2
+    assert select.cache_info().currsize == 0
+
+
 def test_blackwell_launch_table_uses_exact_grid_shape() -> None:
     """Keep production launch pins scoped to the swept device and grid."""
     shape_key = (
@@ -250,6 +402,7 @@ def _compare(
 @pytest.mark.parametrize(
     ("lmax", "n_frames", "n_focus", "channels", "n_grid"), GRID_SHAPES
 )
+@requires_cuda
 def test_float32_matches_eager_conditioning(
     lmax: int, n_frames: int, n_focus: int, channels: int, n_grid: int
 ) -> None:
@@ -260,6 +413,7 @@ def test_float32_matches_eager_conditioning(
 @pytest.mark.parametrize(
     ("lmax", "n_frames", "n_focus", "channels", "n_grid"), GRID_SHAPES
 )
+@requires_cuda
 def test_autocast_bfloat16_matches_eager_conditioning(
     lmax: int, n_frames: int, n_focus: int, channels: int, n_grid: int
 ) -> None:
@@ -267,6 +421,7 @@ def test_autocast_bfloat16_matches_eager_conditioning(
     _compare((lmax, n_frames, n_focus, channels, n_grid), amp=True)
 
 
+@requires_cuda
 def test_noncontiguous_operands_match_eager_conditioning() -> None:
     """Cover the channel-slice strides supplied by production grid nets."""
     _compare(GRID_SHAPES[1], amp=True, strided=True)
@@ -276,6 +431,7 @@ def test_noncontiguous_operands_match_eager_conditioning() -> None:
     ("n_focus", "degree_major"),
     [(1, False), (2, True)],
 )
+@requires_cuda
 def test_symbolic_graph_reuses_layout_across_batch_shapes(
     n_focus: int, degree_major: bool
 ) -> None:
@@ -354,6 +510,7 @@ def evaluate(
         torch.testing.assert_close(got, want)
 
 
+@requires_cuda
 def test_aotautograd_reuses_multifocus_stride_across_batch_shapes() -> None:
     """Keep a producer's batch-dependent coefficient stride out of guards."""
     n_frames, coeff_dim, n_focus, channels, n_grid = 3, 4, 2, 16, 32
@@ -391,3 +548,201 @@ def evaluate(
             torch.testing.assert_close(got, want)
 
     assert compile_counter.frame_count == 1
+
+
+@pytest.fixture
+def tf32_disabled() -> Iterator[None]:
+    """Disable TF32 in cuBLAS and cuDNN for one test and restore it afterwards.
+
+    The flags are restored through their own attributes: the cuBLAS flag is
+    served by the ``__getattr__`` of its module rather than an attribute
+    entry, which ``monkeypatch.setattr`` of pytest 9.2 records as absent and
+    deletes at teardown.
+    """
+    backends = (torch.backends.cuda.matmul, torch.backends.cudnn)
+    previous = [backend.allow_tf32 for backend in backends]
+    for backend in backends:
+        backend.allow_tf32 = False
+    yield
+    for backend, allow_tf32 in zip(backends, previous, strict=True):
+        backend.allow_tf32 = allow_tf32
+
+
+@pytest.mark.usefixtures("tf32_disabled")
+@pytest.mark.parametrize(
+    ("n_focus", "channels", "amp", "force_reference"),
+    [
+        (2, 96, False, False),
+        (1, 256, False, False),
+        (2, 96, True, False),
+        (1, 256, True, False),
+        (2, 16, False, True),
+    ],
+    ids=[
+        "multifocus-fp32",
+        "singlefocus-fp32",
+        "multifocus-amp",
+        "singlefocus-amp",
+        "reference",
+    ],
+)
+@requires_cuda
+def test_cold_compiled_force_backward(
+    n_focus: int,
+    channels: int,
+    amp: bool,
+    force_reference: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Resolve device resources before tracing a production grid's second order.
+
+    The first execution is compiled, so an eager pass cannot populate the
+    launch cache first. The two production layouts cover multi-focus messages
+    and the single-focus feed-forward grid. A device with no viable fused
+    configuration must retain the reference composition through compilation.
+    """
+    # Inductor requires native BF16 support, not autocast's software emulation.
+    if amp and not torch.cuda.is_bf16_supported(including_emulation=False):
+        pytest.skip("Compiled BF16 requires native device support")
+    torch._dynamo.reset()
+    fallback_ops = set()
+    probed_ops = set()
+    admitted_ops = set()
+    probe_inputs = {}
+    reference = grid_pair_module._eager_packed
+
+    def record_fallback(op: str, *args: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        fallback_ops.add(op)
+        return reference(op, *args)
+
+    monkeypatch.setattr(grid_pair_module, "_eager_packed", record_fallback)
+    for op, kernel in (
+        ("fwd", grid_pair_module._grid_pair_fwd_kernel),
+        ("bwd", grid_pair_module._grid_pair_bwd_kernel),
+        ("bwd2", grid_pair_module._grid_pair_bwd2_kernel),
+    ):
+
+        def validate_candidate(
+            *args: object,
+            _op: str = op,
+            _warmup: Callable[..., CompiledKernel] = kernel.warmup,
+            **kwargs: object,
+        ) -> CompiledKernel:
+            probed_ops.add(_op)
+            probe_inputs[_op] = (_warmup, args, kwargs)
+            if force_reference:
+                raise grid_pair_module.triton.runtime.errors.OutOfResources(
+                    2, 1, "test"
+                )
+            candidate = _warmup(*args, **kwargs)
+            candidate._init_handles()
+            admitted_ops.add(_op)
+            return candidate
+
+        monkeypatch.setattr(kernel, "warmup", validate_candidate)
+
+    # Degree six reproduces the P=147 resource pressure. The distinct grid
+    # count prevents eager production-shape tests from warming this launch key.
+    n_frames, coeff_dim, n_grid = 3, 49, 587
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(40529)
+    projectors = tuple(
+        torch.randn(
+            n_grid,
+            coeff_dim * n_frames,
+            device=device,
+            dtype=torch.float32,
+            generator=generator,
+        )
+        / n_grid**0.5
+        for _ in range(2)
+    )
+
+    def make_inputs(n_node: int) -> tuple[torch.Tensor, ...]:
+        # Degree-major storage gives the coefficient axis a batch-dependent
+        # stride; the logical entry remains (N, D, F, K * C).
+        shape = (coeff_dim, n_node, n_focus, n_frames * channels)
+        return tuple(
+            torch.randn(
+                shape,
+                device=device,
+                dtype=torch.float32,
+                generator=generator,
+                requires_grad=True,
+            )
+            for _ in range(2)
+        )
+
+    def force_forward(
+        left: torch.Tensor,
+        right: torch.Tensor,
+        to_grid: torch.Tensor,
+        from_grid: torch.Tensor,
+    ) -> tuple[torch.Tensor, ...]:
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
+            out = grid_pair_train(
+                left.permute(1, 0, 2, 3),
+                right.permute(1, 0, 2, 3),
+                to_grid,
+                from_grid,
+                n_frames,
+            )
+        gradients = torch.autograd.grad(out.sum(), (left, right), create_graph=True)
+        return out, *gradients
+
+    def force_loss_gradient(
+        outputs: tuple[torch.Tensor, ...],
+        inputs: tuple[torch.Tensor, ...],
+    ) -> tuple[torch.Tensor, ...]:
+        loss = sum(value.float().square().mean() for value in outputs)
+        return torch.autograd.grad(loss, inputs)
+
+    traced_inputs = make_inputs(7)
+    graph = make_fx(force_forward, tracing_mode="symbolic")(*traced_inputs, *projectors)
+    counter = CompileCounterWithBackend("inductor")
+    with torch._inductor.config.patch(fx_graph_cache=False):
+        compiled = torch.compile(graph, backend=counter, dynamic=True, fullgraph=True)
+        for n_node in (7, 11):
+            inputs = traced_inputs if n_node == 7 else make_inputs(n_node)
+            actual = compiled(*inputs, *projectors)
+            actual_gradients = force_loss_gradient(actual, inputs)
+            # Observe compiled selection before eager evaluation can add events.
+            # A cached/unconditional rejection without validation cannot pass.
+            assert probed_ops == {"fwd", "bwd", "bwd2"}
+            assert fallback_ops == probed_ops - admitted_ops
+            if force_reference:
+                assert not admitted_ops
+            else:
+                for op in probed_ops - admitted_ops:
+                    warmup, args, kwargs = probe_inputs[op]
+                    # Probe only after the cold compiled execution. This uses
+                    # the original kernel, independently of the selector's
+                    # admission decision, without warming the traced path.
+                    try:
+                        minimum = warmup(
+                            *args,
+                            **(
+                                kwargs
+                                | {
+                                    "C_BLK": 16,
+                                    "BLOCK_G": 16,
+                                    "num_stages": 1,
+                                    "num_warps": 4,
+                                }
+                            ),
+                        )
+                        minimum._init_handles()
+                    except grid_pair_module.triton.runtime.errors.OutOfResources:
+                        continue
+                    pytest.fail(
+                        f"{op} rejected every tile although its minimum tile loads"
+                    )
+            expected = force_forward(*inputs, *projectors)
+            expected_gradients = force_loss_gradient(expected, inputs)
+            for got, want in zip(
+                (*actual, *actual_gradients),
+                (*expected, *expected_gradients),
+                strict=True,
+            ):
+                torch.testing.assert_close(got, want)
+    assert counter.frame_count == 1
diff --git a/source/tests/pt_expt/kernels/test_long_range_sum.py b/source/tests/pt_expt/kernels/test_long_range_sum.py
new file mode 100644
index 0000000000..165292f733
--- /dev/null
+++ b/source/tests/pt_expt/kernels/test_long_range_sum.py
@@ -0,0 +1,551 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Correctness of the operator of the latent-charge long range.
+
+For non-periodic frames the CUDA kernel evaluates every pair in float32 from
+positions relative to the first node of each frame, carried as float32 head
+and tail, and sums in compensated float32 before a float64 reduction; the
+total-charge constraint and the per-frame reductions sum in float64 in a fixed
+order. It is held to the float64 evaluation over the explicit pairs, at a
+tolerance relative to the scale of each output, and to the exact invariants of
+a pair sum: padding nodes and frames do not interact, the forces of a frame
+cancel, and the energy, force, charge gradient and virial follow from one
+another.
+
+For periodic frames the CUDA kernels evaluate the mesh sum in float32 on
+B-spline weights formed in float64, and are held to the float64 mesh
+evaluation of the training path, differentiated by autograd.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    Any,
+)
+
+import numpy as np
+import pytest
+import torch
+
+from deepmd.dpmodel.long_range import (
+    COULOMB_CONSTANT,
+)
+from deepmd.dpmodel.utils.neighbor_graph import (
+    mesh_dims_of_cells,
+)
+from deepmd.pt_expt.kernels.long_range import (
+    _reference,
+    kernel_operator_arguments,
+    latent_charge_sum,
+    op_available,
+)
+from deepmd.pt_expt.long_range.kernel import (
+    LESKernel,
+    SOGKernel,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or not op_available(),
+    reason="the long-range operator is CUDA only",
+)
+
+KERNELS = {
+    "les": lambda: LESKernel(alpha=0.8),
+    "sog": lambda: SOGKernel(
+        amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2], precision="float32"
+    ),
+}
+
+
+def _system(
+    n_node: list[int], channels: int, padding: int, seed: int, offset: float
+) -> tuple[torch.Tensor, ...]:
+    """Frames of random nodes far from the origin, with trailing padding nodes.
+
+    Every seventh node carries no charge, and every frame a total charge.
+    """
+    rng = np.random.default_rng(seed)
+    total = sum(n_node) + padding
+    coord = offset + rng.uniform(0.0, 4.0 * max(n_node) ** (1 / 3), (total, 3))
+    charge = rng.normal(0.0, 0.5, (total, channels))
+    device = torch.device("cuda")
+    node_mask = torch.ones(total, dtype=torch.bool, device=device)
+    node_mask[::7] = False
+    return (
+        torch.tensor(coord, dtype=torch.float64, device=device),
+        torch.tensor(charge, dtype=torch.float32, device=device),
+        node_mask,
+        torch.tensor(n_node, dtype=torch.int64, device=device),
+        torch.tensor(
+            rng.integers(-2, 3, len(n_node)), dtype=torch.float64, device=device
+        ),
+    )
+
+
+def _reference_of(
+    system: tuple[torch.Tensor, ...],
+    kernel: Any,
+    channel: int,
+    box: torch.Tensor | None = None,
+    order: int = 8,
+    spacing: float = 1.0,
+) -> tuple[torch.Tensor, ...]:
+    """The float64 reference of the operator."""
+    coord, charge, node_mask, n_node, total_charge = system
+    code, coefficient = kernel_operator_arguments(kernel)
+    if box is None:
+        box = torch.zeros(n_node.shape[0], 9, dtype=torch.float64)
+    return _reference(
+        coord,
+        charge.double(),
+        node_mask,
+        n_node,
+        total_charge,
+        box,
+        coefficient,
+        torch.zeros(1, dtype=torch.int64),
+        code,
+        channel,
+        COULOMB_CONSTANT,
+        order,
+        spacing,
+    )
+
+
+def _evaluate(
+    system: tuple[torch.Tensor, ...],
+    kernel: Any,
+    channel: int,
+    box: torch.Tensor | None = None,
+    order: int = 8,
+    spacing: float = 1.0,
+) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]:
+    actual = latent_charge_sum(
+        *system,
+        kernel,
+        channel,
+        COULOMB_CONSTANT,
+        box=box,
+        order=order,
+        spacing=spacing,
+    )
+    return actual, _reference_of(system, kernel, channel, box, order, spacing)
+
+
+NAMES = ("atom_energy", "charge_gradient", "force", "atom_virial", "virial")
+
+
+@pytest.mark.parametrize("kernel_name", ["les", "sog"])
+@pytest.mark.parametrize(
+    ("n_node", "channels", "padding", "channel"),
+    [
+        ([37, 1, 0, 300], 1, 5, 0),
+        ([37, 1, 0, 300], 1, 5, -1),
+        ([2048], 3, 0, -1),
+        ([5, 129], 8, 3, 2),
+        # A frame of 47 owner blocks of 24 chunks, which takes the wide tiles
+        # of the pair kernel on devices of up to 3384 resident block slots.
+        ([6000], 2, 4, 1),
+    ],
+)
+def test_matches_float64_pair_sum(
+    kernel_name: str, n_node: list[int], channels: int, padding: int, channel: int
+) -> None:
+    """Every output agrees with the float64 pair sum."""
+    system = _system(n_node, channels, padding, 7, 1.0e3)
+    actual, expected = _evaluate(system, KERNELS[kernel_name](), channel)
+    _assert_close(actual, expected)
+    # Padding nodes join no frame.
+    if padding:
+        for value in actual[:4]:
+            assert torch.all(value[-padding:] == 0)
+
+
+@pytest.mark.parametrize("alpha", [1e-3, 0.03, 0.3])
+@pytest.mark.parametrize("channel", [0, -1])
+def test_les_matches_float64_pair_sum_at_small_alpha_r(
+    alpha: float, channel: int
+) -> None:
+    """The erf kernel keeps its precision where alpha r is small.
+
+    The two terms of the slope of erf(x) / x agree to second order in x, so
+    that the slope of an approximate erf loses precision in proportion to
+    1 / x^2 at atomic distances and a small alpha.
+    """
+    system = _system([96], 2, 0, 5, 0.0)
+    _assert_close(*_evaluate(system, LESKernel(alpha=alpha), channel))
+
+
+def _assert_close(
+    actual: tuple[torch.Tensor, ...],
+    expected: tuple[torch.Tensor, ...],
+    tolerance: float = 2e-5,
+) -> None:
+    """Every output agrees with the float64 reference, relative to its largest entry."""
+    for name, value, reference in zip(NAMES, actual, expected, strict=True):
+        assert value.dtype == (
+            torch.float64 if name == "atom_energy" else torch.float32
+        )
+        scale = reference.abs().max().item()
+        error = (value.double() - reference.double()).abs().max().item()
+        assert error <= tolerance * scale, (name, error, scale)
+
+
+@pytest.mark.parametrize(
+    ("kernel_name", "n_node"),
+    [
+        ("les", [64, 200]),
+        ("sog", [64, 200]),
+        # A first frame that ends inside a block of the pair kernel, whose
+        # chunks then reach past the partners of that frame. The erf kernel
+        # decides its far field over the lanes of a warp, which such a block
+        # shares between the frames, so that it is checked on frames of whole
+        # warps.
+        ("sog", [150, 200]),
+    ],
+)
+def test_invariants(kernel_name: str, n_node: list[int]) -> None:
+    """Frames are independent, forces cancel, and the virial is the pair virial."""
+    system = _system(n_node, 2, 0, 11, 0.0)
+    coord, charge, node_mask, n_node_t, total_charge = system
+    together = latent_charge_sum(
+        coord,
+        charge,
+        node_mask,
+        n_node_t,
+        total_charge,
+        KERNELS[kernel_name](),
+        1,
+        COULOMB_CONSTANT,
+    )
+    _, _, force, atom_virial, virial = together
+    # The first frame alone gives the same result.
+    first = n_node[0]
+    alone = latent_charge_sum(
+        coord[:first],
+        charge[:first],
+        node_mask[:first],
+        n_node_t[:1],
+        total_charge[:1],
+        KERNELS[kernel_name](),
+        1,
+        COULOMB_CONSTANT,
+    )
+    for both, single in zip(together, alone, strict=True):
+        torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0)
+    # Internal forces sum to zero per frame, and the virial of a frame is
+    # sum_i r_i (x) f_i.
+    for frame, (begin, end) in enumerate(((0, first), (first, sum(n_node)))):
+        frame_force = force[begin:end].double()
+        assert frame_force.sum(0).abs().max() <= 1e-5 * frame_force.abs().max()
+        position = coord[begin:end] - coord[begin]
+        expected = torch.einsum("ni,nj->ij", position, frame_force)
+        scale = expected.abs().max().item()
+        for value in (
+            atom_virial[begin:end].double().sum(0).reshape(3, 3),
+            virial[frame].double().reshape(3, 3),
+        ):
+            torch.testing.assert_close(value, expected, rtol=0, atol=2e-5 * scale)
+
+
+@pytest.mark.parametrize("kernel_name", ["les", "sog"])
+@pytest.mark.parametrize("channel", [-1, 1])
+def test_reference_outputs_are_energy_derivatives(
+    kernel_name: str, channel: int
+) -> None:
+    """The float64 force and charge gradient differentiate the energy."""
+    coord, charge, node_mask, n_node, total_charge = _system([9, 6], 2, 2, 3, 0.0)
+    coord = coord.clone().requires_grad_(True)
+    charge = charge.double().requires_grad_(True)
+    atom_energy, charge_gradient, force, _, _ = _reference_of(
+        (coord, charge, node_mask, n_node, total_charge),
+        KERNELS[kernel_name](),
+        channel,
+    )
+    grad_coord, grad_charge = torch.autograd.grad(atom_energy.sum(), (coord, charge))
+    torch.testing.assert_close(force, -grad_coord, rtol=1e-9, atol=1e-9)
+    torch.testing.assert_close(
+        charge_gradient.detach(), grad_charge, rtol=1e-9, atol=1e-9
+    )
+
+
+def test_bitwise_reproducible() -> None:
+    """Two evaluations of the same input agree bit for bit, with or without cells."""
+    system = _system([3000], 1, 0, 5, 0.0)
+    for box in (None, _cells(1, 40.0)):
+        first = latent_charge_sum(
+            *system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box
+        )
+        second = latent_charge_sum(
+            *system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box
+        )
+        for a, b in zip(first, second, strict=True):
+            assert torch.equal(a, b)
+
+
+# ---------------------------------------------------------------------------
+# Periodic frames
+# ---------------------------------------------------------------------------
+
+
+def _cells(frame_count: int, length: float, seed: int = 3) -> torch.Tensor:
+    """Triclinic cells of about ``length`` per side, row-major, one row per frame."""
+    rng = np.random.default_rng(seed)
+    cells = []
+    for _ in range(frame_count):
+        cell = np.diag(length * rng.uniform(0.8, 1.2, 3))
+        cell[1, 0], cell[2, 0], cell[2, 1] = length * rng.uniform(-0.2, 0.2, 3)
+        cells.append(cell.reshape(9))
+    return torch.tensor(np.asarray(cells), dtype=torch.float64)
+
+
+@pytest.mark.parametrize("kernel_name", ["les", "sog"])
+@pytest.mark.parametrize(
+    ("n_node", "channels", "padding", "channel", "order"),
+    [
+        ([37, 1, 300], 1, 5, 0, 8),
+        ([37, 1, 300], 1, 5, -1, 4),
+        ([500], 3, 0, -1, 6),
+        ([5, 129], 8, 3, 2, 8),
+    ],
+)
+def test_mesh_matches_float64_reference(
+    kernel_name: str,
+    n_node: list[int],
+    channels: int,
+    padding: int,
+    channel: int,
+    order: int,
+) -> None:
+    """Every output of a periodic frame agrees with the mesh evaluation of training.
+
+    The positions extend beyond the cell, so that the fractional coordinates
+    wrap.
+    """
+    system = _system(n_node, channels, padding, 7, -5.0)
+    box = _cells(len(n_node), 12.0)
+    actual, expected = _evaluate(
+        system, KERNELS[kernel_name](), channel, box, order, 0.6
+    )
+    _assert_close(actual, expected, 1e-6)
+    if padding:
+        for value in actual[:4]:
+            assert torch.all(value[-padding:] == 0)
+
+
+@pytest.mark.parametrize("kernel_name", ["les", "sog"])
+def test_mesh_virial_is_the_strain_derivative(kernel_name: str) -> None:
+    """The virial of a periodic frame is the strain derivative of its energy.
+
+    Straining the positions and the cell together by ``I + eps`` changes the
+    energy by ``-W : eps``; the atomic virial of the charged nodes sums to W.
+    The central difference takes the float64 mesh energy, whose rounding stays
+    far below that of the float32 outputs of the operator.
+    """
+    system = _system([60], 2, 0, 13, 0.0)
+    coord, charge, node_mask, n_node, total_charge = system
+    box = _cells(1, 10.0)
+    kernel = KERNELS[kernel_name]()
+    _, _, _, atom_virial, virial = latent_charge_sum(
+        *system, kernel, 0, COULOMB_CONSTANT, box=box, spacing=0.5
+    )
+    h = 1e-5
+    finite = np.zeros((3, 3))
+    for a in range(3):
+        for b in range(3):
+            energies = []
+            for sign in (1.0, -1.0):
+                strain = torch.eye(3, dtype=torch.float64)
+                strain[a, b] += sign * h
+                strained = (
+                    coord @ strain.T.to(coord.device),
+                    charge,
+                    node_mask,
+                    n_node,
+                    total_charge,
+                )
+                cell = box.reshape(3, 3) @ strain.T
+                # Both strained cells keep the mesh of the cell.
+                assert mesh_dims_of_cells(cell.reshape(1, 9).numpy(), 0.5) == (
+                    mesh_dims_of_cells(box.numpy(), 0.5)
+                )
+                energies.append(
+                    _reference_of(strained, kernel, 0, cell.reshape(1, 9), 8, 0.5)[
+                        0
+                    ].sum()
+                )
+            finite[a, b] = -(energies[0] - energies[1]).item() / (2 * h)
+    scale = np.abs(finite).max()
+    np.testing.assert_allclose(
+        virial[0].double().cpu().numpy().reshape(3, 3), finite, atol=2e-6 * scale
+    )
+    np.testing.assert_allclose(
+        atom_virial.double().sum(0).cpu().numpy().reshape(3, 3),
+        finite,
+        atol=2e-6 * scale,
+    )
+
+
+@pytest.mark.parametrize("order", [2, 3, 5, 7, 9, 10, 11, 12])
+def test_mesh_orders(order: int) -> None:
+    """Every B-spline order agrees with the mesh evaluation of training.
+
+    The atomic energies of these small cells are a fraction of the self
+    energy, to a few 1e-7 of which the float32 mesh resolves them.
+    """
+    system = _system([40, 23], 2, 3, 5, -5.0)
+    box = _cells(2, 9.0)
+    actual, expected = _evaluate(system, KERNELS["sog"](), 1, box, order, 0.7)
+    _assert_close(actual, expected, 1e-5)
+
+
+def test_mesh_stencil_wider_than_the_axis() -> None:
+    """A stencil wider than an axis of the mesh wraps around it more than once."""
+    system = _system([12], 1, 0, 9, 0.0)
+    box = _cells(1, 3.0)
+    # Four points per axis, against eight of the stencil.
+    assert mesh_dims_of_cells(box.numpy(), 1.0) == (4, 4, 4)
+    actual, expected = _evaluate(system, KERNELS["les"](), 0, box, 8, 1.0)
+    _assert_close(actual, expected, 1e-6)
+
+
+@pytest.mark.parametrize("periodic", [False, True])
+def test_frame_without_charged_nodes(periodic: bool) -> None:
+    """A frame whose nodes carry no charge contributes zeros, the others are unaffected."""
+    n_node = [40, 6, 50]
+    coord, charge, node_mask, n_node_t, total_charge = _system(n_node, 2, 0, 5, 0.0)
+    node_mask[40:46] = False
+    system = (coord, charge, node_mask, n_node_t, total_charge)
+    box = _cells(3, 10.0) if periodic else None
+    actual, expected = _evaluate(system, KERNELS["les"](), 0, box, 8, 0.6)
+    _assert_close(actual, expected, 1e-6)
+    for value in actual[:4]:
+        assert torch.all(value[40:46] == 0)
+    assert torch.all(actual[4][1] == 0)
+
+
+def test_mesh_frames_are_independent() -> None:
+    """Every periodic frame is evaluated on its own mesh, bit for bit."""
+    n_node = [64, 200]
+    system = _system(n_node, 2, 0, 11, 0.0)
+    coord, charge, node_mask, n_node_t, total_charge = system
+    box = _cells(2, 9.0)
+    kernel = KERNELS["sog"]()
+    together = latent_charge_sum(
+        *system, kernel, 1, COULOMB_CONSTANT, box=box, spacing=0.5
+    )
+    first = n_node[0]
+    alone = latent_charge_sum(
+        coord[:first],
+        charge[:first],
+        node_mask[:first],
+        n_node_t[:1],
+        total_charge[:1],
+        kernel,
+        1,
+        COULOMB_CONSTANT,
+        box=box[:1],
+        spacing=0.5,
+    )
+    for both, single in zip(together, alone, strict=True):
+        torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0)
+
+
+def _mesh_sizes(box: torch.Tensor, spacing: float) -> list[tuple[int, int, int]]:
+    """The mesh size of every frame of ``box``."""
+    return [mesh_dims_of_cells(cell.reshape(1, 9).numpy(), spacing) for cell in box]
+
+
+def test_captured_meshes_read_the_cells_at_replay() -> None:
+    """A CUDA graph of periodic frames evaluates the cells of its box at every replay.
+
+    The graph reads the cells where they lie in pinned memory, so that a
+    replay after a strain that keeps the sizes of the meshes agrees bit for
+    bit with an eager evaluation of the strained cells; cells in pageable
+    memory cannot be captured.
+    """
+    system = _system([64, 200], 2, 0, 11, 0.0)
+    kernel = KERNELS["sog"]().to("cuda")
+
+    def evaluate(box: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        return latent_charge_sum(
+            *system, kernel, 1, COULOMB_CONSTANT, box=box, spacing=0.5
+        )
+
+    cells = _cells(2, 9.0)
+    strained = cells * 1.005
+    assert _mesh_sizes(strained, 0.5) == _mesh_sizes(cells, 0.5)
+    expected = {"first": evaluate(cells), "strained": evaluate(strained)}
+    with pytest.raises(RuntimeError, match="pinned"):
+        with torch.cuda.graph(torch.cuda.CUDAGraph()):
+            evaluate(cells)
+    box = cells.pin_memory()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = evaluate(box)
+    for name, cell in (("strained", strained), ("first", cells)):
+        box.copy_(cell)
+        graph.replay()
+        torch.cuda.synchronize()
+        for value, reference in zip(captured, expected[name], strict=True):
+            torch.testing.assert_close(
+                value, reference, rtol=0, atol=0, msg=lambda m, n=name: f"{n}: {m}"
+            )
+
+
+def test_captured_meshes_outlive_other_shapes() -> None:
+    """A CUDA graph of periodic frames replays after 17 other mesh shapes, bit for bit.
+
+    The operator keeps the transforms of 16 mesh shapes and releases the least
+    recently used past them, except those of a live graph, which runs them at
+    every replay.
+    """
+    system = _system([64, 200], 2, 0, 11, 0.0)
+    kernel = KERNELS["sog"]().to("cuda")
+
+    def evaluate(box: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        return latent_charge_sum(
+            *system, kernel, 1, COULOMB_CONSTANT, box=box, spacing=0.5
+        )
+
+    box = _cells(2, 9.0).pin_memory()
+    expected = evaluate(box)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = evaluate(box)
+    others: set[tuple[int, int, int]] = set()
+    length = 9.0
+    while len(others) < 17:
+        length += 0.5
+        cells = _cells(2, length)
+        evaluate(cells)
+        others |= set(_mesh_sizes(cells, 0.5)) - set(_mesh_sizes(box, 0.5))
+    graph.replay()
+    for value, reference in zip(captured, expected, strict=True):
+        torch.testing.assert_close(value, reference, rtol=0, atol=0)
+
+
+def test_a_capture_needs_an_eager_call_for_its_meshes() -> None:
+    """A CUDA graph captures only meshes that an eager call has evaluated."""
+    system = _system([12], 1, 0, 9, 0.0)
+    kernel = KERNELS["les"]().to("cuda")
+    # A cell whose mesh of 16 x 16 x 240 points no other test evaluates.
+    box = torch.tensor(
+        [[4.0, 0.0, 0.0, 0.0, 4.0, 0.0, 0.0, 0.0, 60.0]], dtype=torch.float64
+    ).pin_memory()
+    graph = torch.cuda.CUDAGraph()
+    with pytest.raises(RuntimeError, match="none has evaluated the mesh of 16 x 16"):
+        with torch.cuda.graph(graph):
+            latent_charge_sum(
+                *system, kernel, 0, COULOMB_CONSTANT, box=box, spacing=0.25
+            )
+
+
+def test_mixed_frames_are_rejected() -> None:
+    """A call holds periodic frames or non-periodic ones, not both."""
+    system = _system([10, 10], 1, 0, 1, 0.0)
+    box = _cells(2, 9.0)
+    box[1] = 0.0
+    with pytest.raises(RuntimeError, match="periodic"):
+        latent_charge_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box)
diff --git a/source/tests/pt_expt/kernels/test_so2_edge_cohort.py b/source/tests/pt_expt/kernels/test_so2_edge_cohort.py
new file mode 100644
index 0000000000..33836528d5
--- /dev/null
+++ b/source/tests/pt_expt/kernels/test_so2_edge_cohort.py
@@ -0,0 +1,321 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Generated-input checks for frozen Neo edge-cohort inference."""
+
+from __future__ import (
+    annotations,
+)
+
+from types import (
+    SimpleNamespace,
+)
+
+import pytest
+import torch
+
+pytest.importorskip("triton")
+
+from deepmd.pt_expt.kernels.triton.sezm import so2_edge_cohort as ecs
+from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as sv
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def operands(edges: int, mode: str = "random") -> tuple:
+    """Generate well-conditioned operands without a downloaded checkpoint."""
+    torch.manual_seed(719)
+    u = torch.randn(2, edges, 320, device="cuda") * 0.1
+    if mode == "zero":
+        u.zero_()
+    alpha = torch.rand(edges, 2, device="cuda") * 0.8 + 0.1
+    if mode == "tiny":
+        alpha.fill_(1e-5)
+    w0 = torch.randn(3, 2, 128, 128, device="cuda") * 0.01
+    w1 = torch.randn(3, 2, 192, 192, device="cuda") * 0.01
+    gw = torch.randn(2, 2, 32, 96, device="cuda") * 0.01
+    out, z, uf = sv._mixing_stack_impl(u, alpha, w0, w1, gw, 3, 32, True)
+    go = torch.randn_like(out) * 0.1
+    return (
+        go,
+        out,
+        z,
+        uf,
+        alpha,
+        w0.transpose(2, 3).contiguous(),
+        w1.transpose(2, 3).contiguous(),
+        gw,
+        gw.transpose(2, 3).contiguous(),
+        3,
+        32,
+        True,
+    )
+
+
+@pytest.mark.parametrize("edges", [0, 1, 37, 129, 1025, 8193])
+@pytest.mark.parametrize("cap", [64, 8192])
+@pytest.mark.parametrize("mode", ["random", "zero", "tiny"])
+def test_cohort_matches_native(edges: int, cap: int, mode: str) -> None:
+    """Cover tail cohorts, two focus strides, legal tiny alpha and no mutation."""
+    args = operands(edges, mode)
+    copies = [x.clone() for x in args[:9]]
+    expected_u, expected_a = sv._mixing_stack_bwd_impl(*args)
+    go, out, z, _, alpha, w0t, w1t, gw, gwt, *_ = args
+    actual_u = ecs.cohort_input_backward_impl(go, z, alpha, w0t, w1t, gw, gwt, cap)
+    actual_a = ecs._cohort_alpha_backward(go, out, alpha)
+    for actual, expected in ((actual_u, expected_u), (actual_a, expected_a)):
+        assert torch.isfinite(actual).all()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    for before, after in zip(copies, args[:9], strict=True):
+        assert torch.equal(before, after)
+
+
+def test_compiled_dynamic_cohort() -> None:
+    """The opaque runtime edge loop supports online dynamic graph execution."""
+    compiled = torch.compile(ecs.cohort_backward, fullgraph=True, dynamic=True)
+    for edges in (37, 129):
+        args = operands(edges)
+        actual = compiled(*args)
+        expected = sv._mixing_stack_bwd_impl(*args)
+        for a, b in zip(actual, expected, strict=True):
+            torch.testing.assert_close(a, b, atol=0, rtol=0)
+
+
+def test_make_fx_retains_runtime_cohorts(monkeypatch) -> None:
+    """A small tracing example must not disable batching on larger inputs."""
+    from torch.fx.experimental.proxy_tensor import (
+        make_fx,
+    )
+
+    graph = make_fx(ecs.cohort_backward, tracing_mode="symbolic")(*operands(37))
+    assert any(
+        node.op == "call_function"
+        and node.target == torch.ops.sezm_triton.so2_cohort_input_bwd.default
+        for node in graph.graph.nodes
+    )
+    calls = []
+    original = ecs.cohort_input_backward_impl
+
+    def observed(go, z, alpha, w0t, w1t, gw, gwt, cap):
+        calls.append((go.shape[0], cap))
+        return original(go, z, alpha, w0t, w1t, gw, gwt, cap)
+
+    monkeypatch.setattr(ecs, "cohort_input_backward_impl", observed)
+    compiled = torch.compile(graph, fullgraph=True, dynamic=True)
+    for edges in (37, 8192, 8193, 16384, 16385):
+        args = operands(edges)
+        actual = compiled(*args)
+        expected = sv._mixing_stack_bwd_impl(*args)
+        assert calls[-1] == (edges, ecs.COHORT_EDGES)
+        for got, ref in zip(actual, expected, strict=True):
+            torch.testing.assert_close(got, ref, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("dynamic", [False, True])
+def test_compiled_native_backward(dynamic: bool) -> None:
+    """Optional strides must preserve the default native compilation path."""
+    compiled = torch.compile(sv._mixing_stack_bwd_op, fullgraph=True, dynamic=dynamic)
+    for edges in (37, 129):
+        args = operands(edges)
+        expected = sv._mixing_stack_bwd_impl(*args)
+        actual = compiled(*args)
+        for a, b in zip(actual, expected, strict=True):
+            torch.testing.assert_close(a, b, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize(
+    "lmax,cf,apply", [(2, 32, True), (3, 64, True), (3, 32, False)]
+)
+def test_unsupported_layout(lmax: int, cf: int, apply: bool) -> None:
+    """Untested layouts must not enter the experimental implementation."""
+    args = operands(1)
+    assert not ecs.supports_edge_cohort(args[0], args[2], lmax, cf, apply)
+    assert not ecs.supports_edge_cohort(args[0].cpu(), args[2].cpu(), 3, 32, True)
+
+
+@pytest.mark.parametrize("reason", ["weights", "saved_z", "saved_u", "create_graph"])
+def test_training_dispatch_preserved(monkeypatch, reason: str) -> None:
+    """Weight/state/higher-order requests bypass automatic cohort selection."""
+    args = operands(1)
+    go, out, z, uf, alpha, w0t, w1t, gw, *_ = args
+    u = torch.zeros_like(uf)
+    w0 = w0t.transpose(2, 3).contiguous()
+    w1 = w1t.transpose(2, 3).contiguous()
+    ctx = SimpleNamespace(
+        saved_tensors=(u, alpha, out, z, uf, w0, w1, gw),
+        needs_input_grad=(
+            True,
+            True,
+            reason == "weights",
+            False,
+            False,
+            False,
+            False,
+            False,
+        ),
+        lmax=3,
+        focus_dim=32,
+        apply_alpha=True,
+    )
+    calls = []
+
+    def training(*unused):
+        calls.append(True)
+        return torch.zeros_like(u), torch.zeros_like(alpha), w0, w1, gw
+
+    def forbidden(*unused):
+        raise AssertionError("ECS entered a training/higher-order request")
+
+    monkeypatch.setattr(sv, "_mixing_stack_train_bwd_op", training)
+    monkeypatch.setattr(ecs, "cohort_backward", forbidden)
+    with torch.set_grad_enabled(reason == "create_graph"):
+        sv._mixing_stack_backward(
+            ctx,
+            go,
+            torch.zeros_like(z) if reason == "saved_z" else None,
+            torch.zeros_like(uf) if reason == "saved_u" else None,
+        )
+    assert calls == [True]
+
+
+def test_frozen_dispatch(monkeypatch) -> None:
+    """The production dispatcher must enter ECS, not an equal native fallback."""
+    args = operands(37)
+    go, out, z, uf, alpha, w0t, w1t, gw, *_ = args
+    ctx = SimpleNamespace(
+        saved_tensors=(
+            uf,
+            alpha,
+            out,
+            z,
+            uf,
+            w0t.transpose(2, 3).contiguous(),
+            w1t.transpose(2, 3).contiguous(),
+            gw,
+        ),
+        needs_input_grad=(True, True, False, False, False, False, False, False),
+        lmax=3,
+        focus_dim=32,
+        apply_alpha=True,
+    )
+    expected = sv._mixing_stack_bwd_impl(*args)
+    original = ecs.cohort_backward
+    calls = []
+
+    def observed(*inputs):
+        calls.append(True)
+        return original(*inputs)
+
+    def forbidden(*unused):
+        raise AssertionError("supported inference silently used the native fallback")
+
+    monkeypatch.setattr(ecs, "cohort_backward", observed)
+    monkeypatch.setattr(sv, "_mixing_stack_bwd_op", forbidden)
+    with torch.no_grad():
+        actual = sv._mixing_stack_backward(ctx, go, None, None)
+    assert calls == [True]
+    for a, b in zip(actual[:2], expected, strict=True):
+        torch.testing.assert_close(a, b, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+def test_production_value_path_enters_cohort(monkeypatch, backend):
+    """A real level-2 convolution must use cohorts on frozen first backward."""
+    for name in (
+        "DP_TRITON_TRAIN",
+        "DP_CUDA_INFER",
+        "DP_CUDA_TRAIN",
+        "DP_CUTE_INFER",
+        "DP_CUTILE_INFER",
+    ):
+        monkeypatch.setenv(name, "0")
+    monkeypatch.setenv("DP_TRITON_INFER", "2")
+    if backend == "pt":
+        from deepmd.pt.model.descriptor.sezm_nn.so2 import (
+            SO2Convolution,
+        )
+
+        precision = {"dtype": torch.float32}
+    else:
+        from deepmd.pt_expt.descriptor.dpa4_nn.so2 import (
+            SO2Convolution,
+        )
+
+        precision = {"precision": "float32"}
+    conv = (
+        SO2Convolution(
+            lmax=3,
+            mmax=1,
+            channels=32,
+            n_focus=2,
+            focus_dim=32,
+            mixing_layers=3,
+            radial_so2_mode="degree_channel",
+            radial_so2_rank=1,
+            n_atten_head=1,
+            seed=7,
+            trainable=False,
+            **precision,
+        )
+        .to("cuda")
+        .eval()
+    )
+    x = (torch.randn(7, 16, 64, device="cuda") * 0.1).requires_grad_()
+    radial = torch.randn(17, 4, 32, device="cuda") * 0.1
+    cache = SimpleNamespace(
+        src=torch.arange(17, device="cuda") % 7,
+        D_full=torch.eye(16, device="cuda").repeat(17, 1, 1),
+    )
+    calls = []
+    original = ecs.cohort_backward
+
+    def observe(*args):
+        calls.append(True)
+        return original(*args)
+
+    monkeypatch.setattr(ecs, "cohort_backward", observe)
+    output = conv._triton_value_path(x, cache, radial)[0]
+    gradient = torch.autograd.grad(output.square().sum(), x)[0]
+    assert calls == [True]
+    assert torch.isfinite(gradient).all()
+
+
+def test_short_force_loss_training_matches_reference(monkeypatch):
+    """Three SGD steps preserve loss, gradients and weights on the Neo layout."""
+    torch.manual_seed(719)
+    actual = [
+        (torch.randn(2, 17, 320, device="cuda") * 0.1).requires_grad_(),
+        (torch.rand(17, 2, device="cuda") * 0.8 + 0.1).requires_grad_(),
+        (torch.randn(3, 2, 128, 128, device="cuda") * 0.01).requires_grad_(),
+        (torch.randn(3, 2, 192, 192, device="cuda") * 0.01).requires_grad_(),
+        (torch.randn(2, 2, 32, 96, device="cuda") * 0.01).requires_grad_(),
+    ]
+    reference = [t.detach().clone().requires_grad_() for t in actual]
+    optimizers = [
+        torch.optim.SGD(values[2:], lr=0.01) for values in (actual, reference)
+    ]
+
+    def forbidden(*unused):
+        raise AssertionError("training entered inference-only ECS")
+
+    monkeypatch.setattr(ecs, "cohort_backward", forbidden)
+    for _ in range(3):
+        results = []
+        for operation, values, optimizer in zip(
+            (sv._mixing_stack_op, sv._mixing_stack_reference),
+            (actual, reference),
+            optimizers,
+            strict=True,
+        ):
+            optimizer.zero_grad(set_to_none=True)
+            output = operation(*values, 3, 32, True)[0]
+            energy = output.square().mean()
+            force = torch.autograd.grad(energy, values[0], create_graph=True)[0]
+            loss = energy + 0.01 * force.square().sum()
+            gradients = torch.autograd.grad(loss, values)
+            results.append((loss.detach(), *[g.detach() for g in gradients]))
+            for parameter, gradient in zip(values[2:], gradients[2:], strict=True):
+                parameter.grad = gradient
+            optimizer.step()
+        for got, expected in zip(*results, strict=True):
+            torch.testing.assert_close(got, expected, atol=1e-5, rtol=1e-4)
+        for got, expected in zip(actual[2:], reference[2:], strict=True):
+            torch.testing.assert_close(got, expected, atol=1e-5, rtol=1e-4)
diff --git a/source/tests/pt_expt/kernels/test_so2_edge_cohort_contract.py b/source/tests/pt_expt/kernels/test_so2_edge_cohort_contract.py
new file mode 100644
index 0000000000..7d355a26af
--- /dev/null
+++ b/source/tests/pt_expt/kernels/test_so2_edge_cohort_contract.py
@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Selection and bounded scratch contracts, without executing GPU kernels."""
+
+from types import (
+    SimpleNamespace,
+)
+from unittest.mock import (
+    patch,
+)
+
+import pytest
+import torch
+from torch._subclasses.fake_tensor import (
+    FakeTensorMode,
+)
+
+pytest.importorskip("triton")
+
+from deepmd.pt_expt.kernels.triton.sezm import so2_edge_cohort as ecs
+from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as sv
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+@pytest.mark.parametrize("level", [None, "0", "1", "2", "3"])
+def test_existing_inference_levels(monkeypatch, backend, level):
+    """Real constructors read the existing control before binding the stack."""
+    for name in (
+        "DP_TRITON_TRAIN",
+        "DP_CUDA_INFER",
+        "DP_CUDA_TRAIN",
+        "DP_CUTE_INFER",
+        "DP_CUTILE_INFER",
+    ):
+        monkeypatch.setenv(name, "0")
+    if level is None:
+        monkeypatch.delenv("DP_TRITON_INFER", raising=False)
+    else:
+        monkeypatch.setenv("DP_TRITON_INFER", level)
+    if backend == "pt":
+        from deepmd.pt.model.descriptor.sezm_nn.so2 import (
+            SO2Convolution,
+        )
+
+        precision = {"dtype": torch.float32}
+    else:
+        from deepmd.pt_expt.descriptor.dpa4_nn.so2 import (
+            SO2Convolution,
+        )
+
+        precision = {"precision": "float32"}
+    conv = (
+        SO2Convolution(
+            lmax=3,
+            mmax=1,
+            channels=32,
+            n_focus=2,
+            focus_dim=32,
+            mixing_layers=3,
+            radial_so2_mode="degree_channel",
+            radial_so2_rank=1,
+            n_atten_head=1,
+            seed=7,
+            trainable=False,
+            **precision,
+        )
+        .cpu()
+        .eval()
+    )
+    assert conv.triton_infer_level == int(level or 0)
+    path = conv._triton_value_path
+    if int(level or 0) < 2:
+        assert path is None
+    else:
+        assert isinstance(path, sv._TritonSO2ValuePath)
+        if level == "2":
+            assert path._stack_op is sv._mixing_stack_op
+        elif sv.stack_fp16x3_configs(32, 3) is not None:
+            assert path._stack_op is not sv._mixing_stack_op
+
+
+class _LaunchRecorder:
+    def __init__(self):
+        self.calls = []
+
+    def __getitem__(self, grid):
+        def record(*args, **kwargs):
+            self.calls.append((args, kwargs))
+
+        return record
+
+
+@pytest.mark.parametrize("changed", [1, 5, 6, 7])
+@pytest.mark.parametrize("kind", ["precision", "stride"])
+def test_unsupported_parameters_keep_native(monkeypatch, changed, kind):
+    """FP32 outputs alone do not qualify every parameter dtype or layout."""
+
+    def forbidden(*args):
+        raise AssertionError("mixed-precision inference entered ECS")
+
+    monkeypatch.setattr(ecs, "cohort_backward", forbidden)
+    with FakeTensorMode():
+        go = torch.empty(37, 2, 320, device="cuda")
+        u = torch.empty(2, 37, 320, device="cuda")
+        alpha = torch.empty(37, 2, device="cuda")
+        z = torch.empty(2, 2, 37, 320, device="cuda")
+        w0 = torch.empty(3, 2, 128, 128, device="cuda")
+        w1 = torch.empty(3, 2, 192, 192, device="cuda")
+        gw = torch.empty(2, 2, 32, 96, device="cuda")
+        saved = [u, alpha, u, z, u, w0, w1, gw]
+        if kind == "precision":
+            saved[changed] = saved[changed].to(torch.bfloat16)
+        else:
+            saved[changed] = (
+                saved[changed].transpose(-2, -1).contiguous().transpose(-2, -1)
+            )
+            assert not saved[changed].is_contiguous()
+        ctx = SimpleNamespace(
+            saved_tensors=saved,
+            needs_input_grad=(True, True, False, False, False, False, False, False),
+            lmax=3,
+            focus_dim=32,
+            apply_alpha=True,
+        )
+        with (
+            patch.object(
+                sv,
+                "_mixing_stack_bwd_op",
+                return_value=(torch.empty_like(u), torch.empty_like(alpha)),
+            ) as native,
+            torch.no_grad(),
+        ):
+            sv._mixing_stack_backward(ctx, go, None, None)
+        native.assert_called_once()
+
+
+@pytest.mark.parametrize("inside", [False, True])
+@pytest.mark.parametrize("edges", [0, 37, 8193, 16385])
+def test_scratch_is_cohort_bounded_and_reused(monkeypatch, inside, edges):
+    """Check allocation sizes/identities at live launches using fake CUDA inputs.
+
+    This tests the real Python traversal and eligibility checks, but substitutes
+    launch recorders for kernels. Numerical GPU coverage lives in the companion
+    test file. A full-graph scratch allocation or per-cohort reallocation fails.
+    """
+    gemm, point, recompute = (_LaunchRecorder() for _ in range(3))
+    monkeypatch.setattr(sv, "_stack_gemm_bwd_kernel", gemm)
+    monkeypatch.setattr(sv, "_stack_point_bwd_kernel", point)
+    monkeypatch.setattr(sv, "_stack_recompute_kernel", recompute)
+    monkeypatch.setattr(
+        sv, "stack_fp32_configs", lambda *args: ((32, 64, 32, 4, 2),) * 3
+    )
+    monkeypatch.setattr(
+        sv,
+        "_point_backward_schedule",
+        lambda *args, **kw: (inside, (16, 4, 2), (16, 4, 2)),
+    )
+    with FakeTensorMode():
+        go = torch.empty(edges, 2, 320, device="cuda")
+        z = torch.empty(2, 2, edges, 320, device="cuda")
+        alpha = torch.empty(edges, 2, device="cuda")
+        w0 = torch.empty(3, 2, 128, 128, device="cuda")
+        w1 = torch.empty(3, 2, 192, 192, device="cuda")
+        gw = torch.empty(2, 2, 32, 96, device="cuda")
+        gwt = torch.empty(2, 2, 96, 32, device="cuda")
+        allocated = []
+        original = torch.Tensor.new_empty
+
+        def observe(tensor, shape, *args, **kwargs):
+            result = original(tensor, shape, *args, **kwargs)
+            allocated.append(result)
+            return result
+
+        with (
+            patch.object(torch.Tensor, "new_empty", observe),
+            patch.object(torch, "empty_like", wraps=torch.empty_like) as like,
+        ):
+            output = ecs.cohort_input_backward_impl(
+                go, z, alpha, w0, w1, gw, gwt, ecs.COHORT_EDGES
+            )
+        assert tuple(output.shape) == (2, edges, 320)
+        if not edges:
+            assert [tuple(t.shape) for t in allocated] == [(2, 0, 320)]
+            assert like.call_count == 0
+            assert not gemm.calls
+            return
+        cap = min(edges, ecs.COHORT_EDGES)
+        expected = [(2, edges, 320)]
+        if edges > cap:
+            expected.append((2, cap, 320))
+        expected.append((2, cap, 96))
+        assert [tuple(t.shape) for t in allocated] == expected
+        residual = allocated[1] if edges > cap else output
+        sigmoid = allocated[-1]
+        assert like.call_count == 1
+        assert like.call_args.args[0] is residual
+        chunks = (edges + cap - 1) // cap
+        assert len(point.calls) == chunks * 2
+        assert len(recompute.calls) == (0 if inside else chunks * 2)
+        scratch_ids = {
+            (id(args[0]), id(args[2]), id(args[4])) for args, _ in point.calls
+        }
+        assert len(scratch_ids) == 1
+        assert next(iter(scratch_ids))[:2] == (id(residual), id(sigmoid))
+        assert all(tuple(args[4].shape) == (2, cap, 320) for args, _ in point.calls)
+        assert all(args[8] <= cap for args, _ in point.calls)
diff --git a/source/tests/pt_expt/kernels/test_so2_edge_cohort_export.py b/source/tests/pt_expt/kernels/test_so2_edge_cohort_export.py
new file mode 100644
index 0000000000..c939cbc36e
--- /dev/null
+++ b/source/tests/pt_expt/kernels/test_so2_edge_cohort_export.py
@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""CPU-traced official exports keep the native, deployable backward."""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA target required"
+)
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+def test_export_keeps_native_backward(monkeypatch, tmp_path, backend):
+    """Stop at the real export boundary and inspect the traced backward ops."""
+    from deepmd.pt.entrypoints import (
+        freeze_pt2,
+    )
+    from deepmd.pt.model.model import (
+        get_model,
+    )
+    from deepmd.pt.train.wrapper import (
+        ModelWrapper,
+    )
+    from deepmd.pt_expt.kernels import (
+        autotune,
+    )
+    from deepmd.pt_expt.kernels.triton.sezm import so2_edge_cohort as ecs
+    from deepmd.pt_expt.utils.serialization import (
+        _trace_and_export_impl,
+    )
+
+    for name in (
+        "DP_TRITON_TRAIN",
+        "DP_CUDA_INFER",
+        "DP_CUDA_TRAIN",
+        "DP_CUTE_INFER",
+        "DP_CUTILE_INFER",
+    ):
+        monkeypatch.setenv(name, "0")
+    monkeypatch.setenv("DP_TRITON_INFER", "2")
+    monkeypatch.setattr(freeze_pt2, "_tune_triton_configs", lambda *args: None)
+    monkeypatch.setattr(autotune, "run_autotune", lambda *args: None)
+    params = {
+        "type": "SeZM",
+        "type_map": ["C"],
+        "use_compile": False,
+        "descriptor": {
+            "type": "SeZM",
+            "sel": [32],
+            "rcut": 4.0,
+            "channels": 32,
+            "n_focus": 2,
+            "focus_dim": 32,
+            "n_radial": 4,
+            "radial_mlp": [32],
+            "use_env_seed": True,
+            "l_schedule": [3, 3, 0],
+            "mmax": 1,
+            "so2_layers": 3,
+            "n_atten_head": 1,
+            "ffn_neurons": 32,
+            "ffn_blocks": 1,
+            "use_amp": False,
+            "precision": "float32",
+            "seed": 7,
+        },
+        "fitting_net": {"neuron": [16], "precision": "float32", "seed": 7},
+    }
+    model = get_model(params).cpu().eval()
+
+    class Boundary(Exception):
+        pass
+
+    def forbidden(*args):
+        raise AssertionError("official export selected Python-runtime cohort launchers")
+
+    monkeypatch.setattr(ecs, "cohort_backward", forbidden)
+
+    def inspect(traced, *args, **kwargs):
+        targets = [str(node.target) for node in traced.graph.nodes]
+        assert any("sezm_triton.so2_mixing_stack_bwd." in name for name in targets)
+        assert not any("sezm_triton.so2_cohort_" in name for name in targets)
+        raise Boundary
+
+    monkeypatch.setattr(torch.export, "export", inspect)
+    with pytest.raises(Boundary):
+        if backend == "pt":
+            checkpoint = tmp_path / "model.pt"
+            wrapper = ModelWrapper(model, model_params=params)
+            torch.save({"model": wrapper.state_dict()}, checkpoint)
+            freeze_pt2.freeze_sezm_to_pt2(
+                str(checkpoint),
+                str(tmp_path / "model.pt2"),
+                device=torch.device("cuda:0"),
+                with_comm=False,
+            )
+        else:
+            _trace_and_export_impl(
+                {"model": model.serialize(), "model_def_script": params},
+                lower_kind="graph",
+                do_atomic_virial=True,
+                target_device=torch.device("cuda:0"),
+            )
diff --git a/source/tests/pt_expt/kernels/test_so2_workspace.py b/source/tests/pt_expt/kernels/test_so2_workspace.py
new file mode 100644
index 0000000000..40d31e03cf
--- /dev/null
+++ b/source/tests/pt_expt/kernels/test_so2_workspace.py
@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Owned residual workspace reuse, without mutating saved or caller tensors."""
+
+from unittest.mock import (
+    patch,
+)
+
+import pytest
+import torch
+
+from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or not vp.SO2_VALUE_PATH_TRITON_AVAILABLE,
+    reason="CUDA and the SO2 Triton value path are required",
+)
+
+
+def _inputs(
+    edges: int, gates: int, lmax: int, focus: int, width: int, apply_alpha: bool
+) -> tuple:
+    """Make well-conditioned operands with nonzero residual and matrix terms."""
+    torch.manual_seed(42)
+    row = (3 * lmax + 1) * width
+    m0, m1 = (lmax + 1) * width, 2 * lmax * width
+    initial = torch.randn(focus, edges, row, device="cuda") * 0.1
+    alpha = torch.rand(edges, focus, device="cuda") + 0.2
+    weight0 = torch.randn(gates + 1, focus, m0, m0, device="cuda") * 0.02
+    weight1 = torch.randn(gates + 1, focus, m1, m1, device="cuda") * 0.02
+    gate = torch.randn(gates, focus, width, lmax * width, device="cuda") * 0.02
+    output, saved, final = vp._mixing_stack_impl(
+        initial, alpha, weight0, weight1, gate, lmax, width, apply_alpha
+    )
+    return (
+        torch.randn_like(output),
+        output,
+        saved,
+        final,
+        alpha,
+        weight0.transpose(2, 3).contiguous(),
+        weight1.transpose(2, 3).contiguous(),
+        gate,
+        gate.transpose(2, 3).contiguous(),
+    )
+
+
+@pytest.mark.parametrize(
+    "lmax,focus,width,edges,gates,apply_alpha",
+    [
+        pytest.param(2, 2, 32, 0, 1, True, id="empty"),
+        pytest.param(2, 2, 32, 37, 0, True, id="final-only"),
+        pytest.param(2, 2, 32, 37, 1, False, id="one-gate-no-alpha"),
+        pytest.param(3, 2, 32, 129, 3, True, id="recurrence-partial-tiles"),
+        pytest.param(3, 1, 64, 37, 1, True, id="single-focus"),
+        pytest.param(6, 2, 96, 37, 3, False, id="bmm-non-power-of-two"),
+        pytest.param(6, 3, 128, 129, 1, True, id="bmm-multiple-focuses"),
+    ],
+)
+def test_backward_matches_reference_without_mutation(
+    lmax: int, focus: int, width: int, edges: int, gates: int, apply_alpha: bool
+) -> None:
+    """Cover empty edges, partial tiles, wide channels and residual recurrences."""
+    inputs = _inputs(edges, gates, lmax, focus, width, apply_alpha)
+    snapshots = [value.clone() for value in inputs]
+    actual = vp._mixing_stack_bwd_impl(*inputs, lmax, width, apply_alpha)
+    assert actual[0].shape == (focus, edges, (3 * lmax + 1) * width)
+    assert actual[1].shape == (edges, focus)
+    if edges:
+        expected = vp._mixing_stack_backward_reference(
+            *inputs, None, None, lmax, width, apply_alpha
+        )
+        torch.testing.assert_close(actual[0], expected[0], atol=5e-5, rtol=1e-4)
+        if apply_alpha:
+            torch.testing.assert_close(actual[1], expected[1], atol=5e-5, rtol=1e-4)
+    for value, snapshot in zip(inputs, snapshots, strict=True):
+        torch.testing.assert_close(value, snapshot, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize(
+    "with_weights,keep", [(False, False), (True, False), (False, True), (True, True)]
+)
+def test_reuse_only_without_weights_or_saved_state(
+    with_weights: bool, keep: bool
+) -> None:
+    """Check allocation identity at real GEMM launches, including training paths."""
+    inputs = _inputs(37, 2, 3, 2, 32, True)
+    original_wrap = vp.wrap_triton
+    aliases = []
+
+    class LaunchProbe:
+        def __init__(self, wrapped) -> None:
+            self.wrapped = wrapped
+
+        def __getitem__(self, grid):
+            launch = self.wrapped[grid]
+
+            def record(*args, **kwargs):
+                # Residual is arg1, independent GEMM input arg0, output arg5.
+                aliases.append(
+                    (
+                        args[1].data_ptr() == args[5].data_ptr(),
+                        args[0].data_ptr() == args[5].data_ptr(),
+                    )
+                )
+                return launch(*args, **kwargs)
+
+            return record
+
+    def observed_wrap(kernel):
+        wrapped = original_wrap(kernel)
+        return LaunchProbe(wrapped) if kernel is vp._stack_gemm_bwd_kernel else wrapped
+
+    with patch.object(vp, "wrap_triton", observed_wrap):
+        actual = vp._stack_backward_traversal(
+            *inputs,
+            None,
+            None,
+            3,
+            32,
+            True,
+            with_weights=with_weights,
+            keep=keep,
+        )
+    assert len(aliases) == 3
+    assert aliases[0] == (False, False)
+    assert aliases[1:] == [(not with_weights and not keep, False)] * 2
+    assert (actual[2] is not None) is with_weights
+    assert (actual[3] is not None) is keep
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("apply_alpha", [False, True])
+def test_reuse_matches_separate_workspace(
+    dtype: torch.dtype, apply_alpha: bool
+) -> None:
+    """Reuse must not change arithmetic in any supported working precision."""
+    inputs = tuple(value.to(dtype) for value in _inputs(129, 3, 3, 2, 32, apply_alpha))
+    snapshots = [value.clone() for value in inputs]
+    actual = vp._stack_backward_traversal(
+        *inputs, None, None, 3, 32, apply_alpha, with_weights=False, keep=False
+    )
+    # Keeping per-layer state uses distinct buffers and identical GEMM arithmetic.
+    expected = vp._stack_backward_traversal(
+        *inputs, None, None, 3, 32, apply_alpha, with_weights=False, keep=True
+    )
+    count = 2 if apply_alpha else 1
+    for result, reference in zip(actual[:count], expected[:count], strict=True):
+        torch.testing.assert_close(result, reference, atol=0, rtol=0)
+    for value, snapshot in zip(inputs, snapshots, strict=True):
+        torch.testing.assert_close(value, snapshot, atol=0, rtol=0)
+
+
+def _differentiable_inputs(trainable_weights: tuple[bool, bool, bool]) -> list:
+    """Small direct-stack operands shared by dispatch and frozen-force checks."""
+    torch.manual_seed(19)
+    lmax, width = 2, 32
+    edges, focus, gates = 7, 2, 2
+    row, m0, m1 = (3 * lmax + 1) * width, (lmax + 1) * width, 2 * lmax * width
+    values = (
+        torch.randn(focus, edges, row, device="cuda") * 0.1,
+        torch.rand(edges, focus, device="cuda") + 0.2,
+        torch.randn(gates + 1, focus, m0, m0, device="cuda") * 0.02,
+        torch.randn(gates + 1, focus, m1, m1, device="cuda") * 0.02,
+        torch.randn(gates, focus, width, lmax * width, device="cuda") * 0.02,
+    )
+    requires = (True, True, *trainable_weights)
+    return [
+        value.clone().requires_grad_(enabled)
+        for value, enabled in zip(values, requires, strict=True)
+    ]
+
+
+@pytest.mark.parametrize(
+    "trainable_weights",
+    [(True, True, True), (True, False, False), (False, False, False)],
+    ids=["trainable", "partially-frozen", "frozen"],
+)
+@pytest.mark.parametrize("create_graph", [False, True])
+def test_backward_dispatch_protects_training(
+    trainable_weights: tuple[bool, bool, bool], create_graph: bool
+) -> None:
+    """Check real autograd routing without executing a training traversal.
+
+    Mocked gradient operators test selection only, not numerical correctness.
+    Existing model training tests cover ordinary trainable weights; the frozen
+    direct-stack numerical gap is covered by the single force-loss check below.
+    """
+    operands = _differentiable_inputs(trainable_weights)
+    output = vp._mixing_stack_op(*operands, 2, 32, True)[0]
+    gradients = tuple(torch.zeros_like(value) for value in operands)
+    with (
+        patch.object(
+            vp, "_mixing_stack_train_bwd_op", return_value=gradients
+        ) as training,
+        patch.object(
+            vp, "_mixing_stack_bwd_op", return_value=gradients[:2]
+        ) as inference,
+    ):
+        torch.autograd.grad(
+            output.sum(),
+            [v for v in operands if v.requires_grad],
+            create_graph=create_graph,
+        )
+    protected = any(trainable_weights) or create_graph
+    assert training.call_count == int(protected)
+    assert inference.call_count == int(not protected)
+
+
+@pytest.mark.parametrize(
+    "saved_output", [1, 2], ids=["pre-activation", "final-activation"]
+)
+def test_saved_output_cotangent_selects_training(saved_output: int) -> None:
+    """Frozen operands still require protected handling of saved-output cotangents."""
+    operands = _differentiable_inputs((False, False, False))
+    output = vp._mixing_stack_op(*operands, 2, 32, True)[saved_output]
+    gradients = tuple(torch.zeros_like(value) for value in operands)
+    with (
+        patch.object(
+            vp, "_mixing_stack_train_bwd_op", return_value=gradients
+        ) as training,
+        patch.object(
+            vp, "_mixing_stack_bwd_op", return_value=gradients[:2]
+        ) as inference,
+    ):
+        torch.autograd.grad(output.sum(), operands[0])
+    training.assert_called_once()
+    inference.assert_not_called()
+    args = training.call_args.args
+    assert args[10 if saved_output == 1 else 11] is not None
+
+
+def test_frozen_force_loss_second_derivative() -> None:
+    """Retain the unique frozen-weight direct-stack force-loss numerical path.
+
+    One small case checks first and second derivatives without repeating
+    optimizer updates or crossing all freezing modes with both gate widths.
+    """
+    actual_inputs = _differentiable_inputs((False, False, False))
+    reference_inputs = [
+        value.detach().clone().requires_grad_(value.requires_grad)
+        for value in actual_inputs
+    ]
+    results = []
+    for operation, operands in (
+        (vp._mixing_stack_op, actual_inputs),
+        (vp._mixing_stack_reference, reference_inputs),
+    ):
+        output = operation(*operands, 2, 32, True)[0]
+        energy = output.square().sum() / 7
+        force = torch.autograd.grad(energy, operands[0], create_graph=True)[0]
+        objective = energy + force.square().sum()
+        gradients = torch.autograd.grad(objective, operands[:2])
+        results.append((output, force, *gradients))
+    for actual, expected in zip(*results, strict=True):
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-4)
diff --git a/source/tests/pt_expt/model/test_dpa4_interop.py b/source/tests/pt_expt/model/test_dpa4_interop.py
index 145eaf3bd6..554fc639ee 100644
--- a/source/tests/pt_expt/model/test_dpa4_interop.py
+++ b/source/tests/pt_expt/model/test_dpa4_interop.py
@@ -168,9 +168,7 @@ def test_adam_routing(
                 bridging_method="ZBL", bridging_r_inner=0.8, bridging_r_outer=1.2
             )
         elif layout == "spin":
-            config["type"] = "standard"
-            config["fitting_net"] = hybrid["fitting_net"]
-            config["spin"] = {"use_spin": [True, False], "virtual_scale": [0.3]}
+            config["spin"] = {"use_spin": [True, False], "scheme": "native"}
         elif layout == "default":
             config = hybrid
             config["descriptor"] = config["descriptor"]["list"][0]
diff --git a/source/tests/pt_expt/model/test_dpa4_spin_contract.py b/source/tests/pt_expt/model/test_dpa4_spin_contract.py
new file mode 100644
index 0000000000..76daa302f7
--- /dev/null
+++ b/source/tests/pt_expt/model/test_dpa4_spin_contract.py
@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""DPA4 spin construction contracts for the exportable backend."""
+
+import copy
+
+import pytest
+
+from deepmd.pt_expt.model.dp_linear_model import (
+    LinearEnergyModel,
+)
+from deepmd.pt_expt.model.get_model import (
+    get_model,
+    get_standard_model,
+)
+from deepmd.pt_expt.model.native_spin_model import (
+    NativeSpinEnergyModel,
+)
+from deepmd.pt_expt.model.spin_ener_model import (
+    SpinEnergyModel,
+)
+from deepmd.utils.spin import (
+    Spin,
+)
+
+
+def _config() -> dict:
+    return {
+        "type_map": ["Ni", "O"],
+        "descriptor": {
+            "type": "dpa4",
+            "rcut": 4.0,
+            "sel": 8,
+            "channels": 4,
+            "n_radial": 3,
+            "lmax": 1,
+            "mmax": 1,
+            "n_blocks": 1,
+            "precision": "float64",
+            "seed": 7,
+            "random_gamma": False,
+        },
+        "fitting_net": {"type": "dpa4_ener", "neuron": [4]},
+    }
+
+
+@pytest.mark.parametrize("descriptor_type", ["dpa4", "DPA4", "sezm", "SeZM"])
+@pytest.mark.parametrize("model_type", [None, "standard"])
+@pytest.mark.parametrize("scheme", [None, "invalid"])
+def test_standard_spin_requires_compatible_backbone(
+    descriptor_type: str, model_type: str | None, scheme: str | None
+) -> None:
+    config = _config()
+    config["descriptor"]["type"] = descriptor_type
+    if model_type is not None:
+        config["type"] = model_type
+    config["spin"] = {"use_spin": [True, False], "virtual_scale": 0.3}
+    if scheme is not None:
+        config["spin"]["scheme"] = scheme
+    with pytest.raises(ValueError, match=r"require.*native"):
+        get_model(config)
+
+
+def test_standard_native_spin_construction() -> None:
+    config = _config()
+    config["spin"] = {"use_spin": [True, False], "scheme": "native"}
+    original = copy.deepcopy(config)
+    model = get_model(config)
+    assert isinstance(model, NativeSpinEnergyModel)
+    assert model.get_type_map() == ["Ni", "O"]
+    assert model.get_sel() == [config["descriptor"]["sel"]]
+    assert model.get_nnei() == config["descriptor"]["sel"]
+    assert config == original
+
+
+@pytest.mark.parametrize(
+    "serialized, depth", [(False, 0), (False, 1), (False, 2), (True, 0)]
+)
+def test_spin_wrapper_requires_compatible_backbone(
+    serialized: bool, depth: int
+) -> None:
+    backbone = get_standard_model(_config())
+    for _ in range(depth):
+        backbone = LinearEnergyModel(
+            models=[backbone.atomic_model], type_map=backbone.get_type_map()
+        )
+    spin = Spin(use_spin=[True, False], virtual_scale=0.3)
+    with pytest.raises(ValueError, match="require a native spin model"):
+        if serialized:
+            SpinEnergyModel.deserialize(
+                {
+                    "type": "spin_ener",
+                    "backbone_model": backbone.serialize(),
+                    "spin": spin.serialize(),
+                }
+            )
+        else:
+            SpinEnergyModel(backbone_model=backbone, spin=spin)
+
+
+@pytest.mark.parametrize("serialized", [False, True])
+def test_spin_wrapper_checks_hybrid_descriptors(serialized: bool) -> None:
+    config = _config()
+    config["descriptor"] = {"type": "hybrid", "list": [config["descriptor"]]}
+    backbone = get_standard_model(config)
+    spin = Spin(use_spin=[True, False], virtual_scale=0.3)
+    with pytest.raises(ValueError, match="require a native spin model"):
+        if serialized:
+            SpinEnergyModel.deserialize(
+                {
+                    "type": "spin_ener",
+                    "backbone_model": backbone.serialize(),
+                    "spin": spin.serialize(),
+                }
+            )
+        else:
+            SpinEnergyModel(backbone_model=backbone, spin=spin)
diff --git a/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py b/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py
index 3a4ee60793..cd26a57c4e 100644
--- a/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py
+++ b/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py
@@ -209,6 +209,25 @@ def test_fold_reproduces_the_referenced_model() -> None:
     )
 
 
+def test_fold_vacuum_reference_keeps_cpu_residency(monkeypatch) -> None:
+    """Folding on CPU must leave every parameter and buffer on CPU.
+
+    Export moves the model to the CPU tracing device before folding. The fold
+    replaces ``bias_atom_e`` (and may set ``vacuum_table``) through numpy
+    assignment; those buffers must not jump back to ``env.DEVICE``.
+    """
+    model = make_model(True).to("cpu")
+    monkeypatch.setattr(env, "DEVICE", torch.device("meta"))
+    model.fold_vacuum_reference()
+    assert not model.atomic_model.fitting_net.vacuum_ref
+    for name, tensor in model.named_parameters():
+        assert tensor.device.type == "cpu", name
+    for name, tensor in model.named_buffers():
+        if tensor is None:
+            continue
+        assert tensor.device.type == "cpu", name
+
+
 def test_dense_and_graph_routes_share_the_vacuum_descriptor() -> None:
     """The reference rows appended to a graph equal the single-atom-frame descriptor."""
     model = make_model(vacuum_ref=True)
@@ -320,6 +339,32 @@ def test_freeze_folds_the_reference(tmp_path, suffix) -> None:
         atol=1e-8,
     )
 
+    if suffix == ".pt2":
+        from deepmd.pt_expt.model.model import (
+            BaseModel,
+        )
+        from deepmd.pt_expt.utils.serialization import (
+            deserialize_to_file,
+            serialize_from_file,
+        )
+
+        data = serialize_from_file(str(tmp_path / "vacuum_frozen.pt2"))
+        restored = BaseModel.deserialize(data["model"])
+        fitting = restored.atomic_model.fitting_net
+        assert fitting.needs_vacuum_descriptor()
+        np.testing.assert_allclose(fitting.bias_atom_e.detach().cpu().numpy(), BIAS)
+        again = tmp_path / "vacuum_again.pt2"
+        deserialize_to_file(str(again), data)
+        dp_again = DeepPot(str(again))
+        for itype in range(2):
+            for evaluator in (dp, dp_again):
+                _, _, _, atom_energy, _ = evaluator.eval(
+                    np.zeros((1, 1, 3)), None, np.array([itype]), atomic=True
+                )
+                np.testing.assert_allclose(
+                    atom_energy.reshape(-1), BIAS[itype], rtol=1e-8, atol=1e-8
+                )
+
 
 def test_freeze_keeps_the_reference_with_fparam(tmp_path) -> None:
     """With frame parameters the frozen model references from the stored vacuum table."""
@@ -388,35 +433,6 @@ def tensor(array: np.ndarray, dtype: torch.dtype = torch.float64) -> torch.Tenso
     assert torch.all(padded[atype_flat.shape[0] :] == 0.0)
 
 
-def test_archive_carries_the_live_model(tmp_path) -> None:
-    """The archive keeps the unfolded model, so a re-export resolves the reference once."""
-    from deepmd.pt_expt.model.model import (
-        BaseModel,
-    )
-    from deepmd.pt_expt.utils.serialization import (
-        deserialize_to_file,
-        serialize_from_file,
-    )
-
-    _, dp = freeze_model(tmp_path, numb_fparam=0)
-    data = serialize_from_file(str(tmp_path / "vacuum_frozen.pt2"))
-    model = BaseModel.deserialize(data["model"])
-    fitting = model.atomic_model.fitting_net
-    assert fitting.needs_vacuum_descriptor()
-    np.testing.assert_allclose(fitting.bias_atom_e.detach().cpu().numpy(), BIAS)
-    again = tmp_path / "vacuum_again.pt2"
-    deserialize_to_file(str(again), data)
-    dp_again = DeepPot(str(again))
-    for itype in range(2):
-        for evaluator in (dp, dp_again):
-            _, _, _, atom_energy, _ = evaluator.eval(
-                np.zeros((1, 1, 3)), None, np.array([itype]), atomic=True
-            )
-            np.testing.assert_allclose(
-                atom_energy.reshape(-1), BIAS[itype], rtol=1e-8, atol=1e-8
-            )
-
-
 def test_without_preset_the_output_is_not_referenced() -> None:
     """A bias fitted from the data is no isolated-atom energy, so the output stays plain."""
     rng = np.random.default_rng(1)
diff --git a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py
index 02af9d412d..ba377697db 100644
--- a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py
+++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py
@@ -14,6 +14,9 @@
 from deepmd.pt_expt.utils import (
     env,
 )
+from deepmd.pt_expt.utils.canonical_graph import (
+    canonical_index_bytes,
+)
 from deepmd.pt_expt.utils.serialization import (
     _resolve_lower_kind,
     _trace_and_export,
@@ -114,6 +117,46 @@ def _run_graph(
     )
 
 
+def _sample(model: torch.nn.Module) -> tuple:
+    """Return a canonical graph, its flat types and a per-node moment."""
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        attach_edge_csr,
+        build_neighbor_graph,
+    )
+
+    nodes = 24
+    generator = torch.Generator(device=env.DEVICE).manual_seed(31)
+    coord = 5.0 * torch.rand(
+        1, nodes, 3, dtype=torch.float32, device=env.DEVICE, generator=generator
+    )
+    atype = torch.arange(nodes, device=env.DEVICE).reshape(1, -1) % 2
+    graph = attach_edge_csr(
+        build_neighbor_graph(coord, atype, None, model.get_rcut()), nodes
+    )
+    spin = torch.randn(
+        nodes, 3, dtype=torch.float32, device=env.DEVICE, generator=generator
+    )
+    return graph, atype.reshape(-1), spin
+
+
+def _assert_exported_outputs(
+    exported: dict[str, torch.Tensor], reference: dict[str, torch.Tensor]
+) -> None:
+    """Compare the outputs of an exported lower with those of the autograd lower."""
+    for key, name in (
+        ("energy", "energy_redu"),
+        ("force", "energy_derv_r"),
+        ("virial", "energy_derv_c_redu"),
+    ):
+        torch.testing.assert_close(
+            exported[key].reshape(-1),
+            reference[name].reshape(-1),
+            atol=8e-6,
+            rtol=1e-4,
+            check_dtype=False,
+        )
+
+
 def test_graph_lower_energy_force_are_finite() -> None:
     model = get_model(_config()).to(env.DEVICE).eval()
     assert model.get_descriptor().uses_graph_lower()
@@ -579,7 +622,87 @@ def test_compact_canonical_graph_export(
     assert isinstance(exported, torch.export.ExportedProgram)
     assert metadata["lower_input_kind"] == "dpa4c_canonical"
     assert metadata["graph_edge_dtype"] == "float32"
-    assert metadata["canonical_index_dtype"] == "uint32"
+    assert metadata["canonical_index_dtype"] == "auto"
+    assert metadata["canonical_index_storage"] == "byte_words"
+
+
+@_GPU
+@pytest.mark.parametrize("channels", [8, 64])  # compressed channel width
+def test_compact_canonical_export_accepts_source_csr_placeholders(
+    monkeypatch: pytest.MonkeyPatch,
+    channels: int,
+) -> None:
+    """The exported compact lower takes placeholders for the source CSR.
+
+    The source CSR keeps extents of its own in the export, so a caller that
+    declines the per-atom virial passes two-entry placeholders in its place.
+    The program then returns a zero per-atom virial and every other output of
+    the complete graph unchanged.
+    """
+    from deepmd.pt_expt.utils.serialization import (
+        build_synthetic_canonical_graph_inputs,
+    )
+
+    monkeypatch.setenv("DP_CUDA_INFER", "2")
+    model = get_model(_compressed_config(channels)).to("cpu").eval()
+    model.get_descriptor().enable_compression(min_nbor_dist=0.5)
+    exported, _metadata, _model_json, _output_keys = _trace_and_export(
+        {"model": model.serialize()},
+        lower_kind="dpa4c_canonical",
+        do_atomic_virial=True,
+    )
+    # Storage past the real edges holds guard rows, so a generous edge axis
+    # keeps every edge of the synthetic system. The inputs reside on the
+    # device the program was exported for.
+    device = next(
+        node.meta["val"].device
+        for node in exported.graph_module.graph.nodes
+        if node.op == "placeholder"
+    )
+    inputs = build_synthetic_canonical_graph_inputs(model, 4096, device=device)
+    program = exported.module()
+    complete = program(*inputs)
+    declined = program(*inputs[:6], inputs[6].new_zeros(2), inputs[7].new_zeros(2))
+    assert complete["atom_virial"].any()
+    assert not declined["atom_virial"].any()
+    assert declined.keys() == complete.keys()
+    for key, value in complete.items():
+        if key != "atom_virial":
+            assert torch.equal(declined[key], value), key
+
+
+@_GPU
+@pytest.mark.parametrize("index_dtype", [torch.uint32, torch.int64])
+def test_canonical_atom_virial_sums_to_frame(
+    monkeypatch: pytest.MonkeyPatch,
+    index_dtype: torch.dtype,
+) -> None:
+    """The complete source CSR attributes the frame virial to its atoms."""
+    monkeypatch.setenv("DP_CUDA_INFER", "2")
+    model = get_model(_compressed_config(64)).to(env.DEVICE).eval()
+    graph, atype, _spin = _sample(model)
+    model.get_descriptor().enable_compression(min_nbor_dist=0.5)
+    physical = int(graph.destination_row_ptr[-1])
+    result = model.forward_lower_canonical_graph(
+        atype,
+        graph.n_node,
+        graph.n_node.clone(),
+        graph.edge_index[0, :physical].to(index_dtype).contiguous(),
+        graph.edge_vec[:physical].contiguous(),
+        graph.destination_row_ptr,
+        graph.source_row_ptr,
+        graph.source_order[:physical].to(torch.int64).contiguous(),
+        do_atomic_virial=True,
+    )
+    assert torch.isfinite(result["atom_virial"]).all()
+    assert result["atom_virial"].any()
+    torch.testing.assert_close(
+        result["atom_virial"].reshape(-1, 3, 3).sum(0),
+        result["virial"].reshape(3, 3),
+        atol=8e-6,
+        rtol=1e-4,
+        check_dtype=False,
+    )
 
 
 @_GPU
@@ -591,6 +714,95 @@ def test_auto_lower_kind_selects_compact_canonical(channels: int) -> None:
     assert _resolve_lower_kind("model.pt2", data, "auto") == "dpa4c_canonical"
 
 
+def test_a_case_embedding_reaches_the_fused_fitting_operator(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A branch of a shared fitting network keeps the fused fitting operator.
+
+    The branch carries a fixed one-hot case embedding, which the fused fitting
+    operator does not read. Export folds it into the first layer of the
+    fitting network, so the exported graph lower evaluates the network with
+    that operator and reproduces the autograd lower of the model as trained.
+    """
+    monkeypatch.setenv("DP_CUDA_INFER", "1")
+    # The fitting network of the compressed configuration is the one that the
+    # fused operator evaluates; the descriptor stays uncompressed.
+    config = _compressed_config()
+    config["fitting_net"]["dim_case_embd"] = 3
+    model = get_model(config).to(env.DEVICE).eval()
+    model.get_fitting_net().set_case_embd(1)
+    exported, metadata, _model_json, _output_keys = _trace_and_export(
+        {"model": model.serialize()},
+        lower_kind="graph",
+        do_atomic_virial=True,
+    )
+    targets = {str(node.target) for node in exported.graph_module.graph.nodes}
+    assert "deepmd.graph_fitting.default" in targets
+    dtype = getattr(torch, metadata["graph_edge_dtype"])
+    sample = build_synthetic_graph_inputs(
+        model, e_max=None, nframes=2, nloc=7, dtype=dtype, device=env.DEVICE
+    )
+    _assert_exported_outputs(exported.module()(*sample), _run_graph(model, dtype))
+
+
+@_GPU
+def test_a_case_embedding_deploys_on_the_compact_canonical_lower(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A branch of a shared fitting network exports like a single-task model.
+
+    The branch carries a fixed one-hot case embedding, which export folds into
+    the first layer of the fitting network, next to the vacuum reference. The
+    compressed model then selects the compact canonical schema, and the
+    exported program reproduces the autograd lower of the model as trained.
+    """
+    monkeypatch.setenv("DP_CUDA_INFER", "2")
+    config = _compressed_config()
+    config["type_map"] = ["O", "H"]
+    config["fitting_net"].update({"dim_case_embd": 3, "vacuum_ref": True})
+    config["preset_out_bias"] = {"energy": {"O": -3.0, "H": 0.5}}
+    model = get_model(config).to(env.DEVICE).eval()
+    model.get_fitting_net().set_case_embd(1)
+    model.get_descriptor().enable_compression(min_nbor_dist=0.5)
+    data = {"model": model.serialize()}
+    assert _resolve_lower_kind("model.pt2", data, "auto") == "dpa4c_canonical"
+    exported, _metadata, _model_json, _output_keys = _trace_and_export(
+        data,
+        lower_kind="dpa4c_canonical",
+        do_atomic_virial=True,
+    )
+
+    graph, atype, _moment = _sample(model)
+    reference = model.forward_common_lower_graph(
+        atype,
+        graph.n_node,
+        graph.n_node.clone(),
+        graph.edge_index,
+        graph.edge_vec,
+        graph.edge_mask,
+        destination_order=graph.destination_order,
+        destination_row_ptr=graph.destination_row_ptr,
+        source_order=graph.source_order,
+        source_row_ptr=graph.source_row_ptr,
+        destination_sorted=bool(graph.destination_sorted),
+    )
+    physical = int(graph.destination_row_ptr[-1])
+    inputs = (
+        atype,
+        graph.n_node,
+        graph.n_node.clone(),
+        canonical_index_bytes(
+            graph.edge_index[0][:physical].to(torch.uint32).contiguous()
+        ),
+        graph.edge_vec[:physical].contiguous(),
+        graph.destination_row_ptr,
+        graph.source_row_ptr,
+        graph.source_order[:physical].to(torch.int64).contiguous(),
+    )
+    deployed = exported.module()(*inputs)
+    _assert_exported_outputs(deployed, reference)
+
+
 def test_the_graph_lower_conditions_each_frame_on_its_own_charge_state() -> None:
     """The frame condition must survive the model seam and stay per frame.
 
@@ -750,28 +962,6 @@ def _spin_config(channels: int = 16) -> dict:
     return config
 
 
-def _spin_sample(model: torch.nn.Module) -> tuple:
-    """Return a canonical graph, its flat types and a per-node moment."""
-    from deepmd.dpmodel.utils.neighbor_graph import (
-        attach_edge_csr,
-        build_neighbor_graph,
-    )
-
-    nodes = 24
-    generator = torch.Generator(device=env.DEVICE).manual_seed(31)
-    coord = 5.0 * torch.rand(
-        1, nodes, 3, dtype=torch.float32, device=env.DEVICE, generator=generator
-    )
-    atype = torch.arange(nodes, device=env.DEVICE).reshape(1, -1) % 2
-    graph = attach_edge_csr(
-        build_neighbor_graph(coord, atype, None, model.get_rcut()), nodes
-    )
-    spin = torch.randn(
-        nodes, 3, dtype=torch.float32, device=env.DEVICE, generator=generator
-    )
-    return graph, atype.reshape(-1), spin
-
-
 @_GPU
 @pytest.mark.parametrize("gate", [0.8, 0.0])  # active vs inactive spin gate
 def test_compressed_spin_lowers_match_autograd(
@@ -800,7 +990,7 @@ def test_compressed_spin_lowers_match_autograd(
     descriptor.spin.set_spin_reference(np.array([1.7, 1.0, 1.0]))
     with torch.no_grad():
         descriptor.spin.spin_gate.fill_(gate)
-    graph, atype, spin = _spin_sample(model)
+    graph, atype, spin = _sample(model)
     descriptor.enable_compression(min_nbor_dist=0.5)
 
     def lower() -> dict[str, torch.Tensor]:
@@ -832,23 +1022,26 @@ def lower() -> dict[str, torch.Tensor]:
         atype,
         graph.n_node,
         graph.n_node.clone(),
-        graph.edge_index[0][:physical].to(torch.uint32).contiguous(),
+        graph.edge_index[0][:physical].to(torch.int64).contiguous(),
         graph.edge_vec[:physical].contiguous(),
         graph.destination_row_ptr,
         graph.source_row_ptr,
-        graph.source_order[:physical].to(torch.uint32).contiguous(),
+        graph.source_order[:physical].to(torch.int64).contiguous(),
         do_atomic_virial=False,
         spin=spin,
     )
+    # The canonical operator returns its fixed-point force sums in fp64.
     torch.testing.assert_close(
         canonical["force_mag"].reshape(-1, 3),
         reference["energy_derv_r_mag"].reshape(-1, 3),
         atol=8e-6,
         rtol=1e-4,
+        check_dtype=False,
     )
     torch.testing.assert_close(
         canonical["force"].reshape(-1, 3),
         reference["energy_derv_r"].reshape(-1, 3),
         atol=8e-6,
         rtol=1e-4,
+        check_dtype=False,
     )
diff --git a/source/tests/pt_expt/model/test_dpa4c_zbl_bridging.py b/source/tests/pt_expt/model/test_dpa4c_zbl_bridging.py
index 3e881fe363..375b1e92a1 100644
--- a/source/tests/pt_expt/model/test_dpa4c_zbl_bridging.py
+++ b/source/tests/pt_expt/model/test_dpa4c_zbl_bridging.py
@@ -544,7 +544,7 @@ def test_canonical_route_matches_the_autograd_composition(
             edge_vec[:physical].contiguous(),
             destination_row_ptr,
             source_row_ptr,
-            source_order[:physical].to(torch.uint32).contiguous(),
+            source_order[:physical].to(torch.int64).contiguous(),
             do_atomic_virial=True,
         )
         for key, reference_key in (
diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py
index ecead9a672..a73674e48a 100644
--- a/source/tests/pt_expt/model/test_get_model_bridging.py
+++ b/source/tests/pt_expt/model/test_get_model_bridging.py
@@ -185,7 +185,7 @@ def test_fine_tune_locates_the_learned_part_by_structure() -> None:
     A descriptor exclusion gives the child its own pair exclusion, which keeps
     the composition off the fused route, but not its learned part off its path.
     """
-    from deepmd.pt_expt.train.training import (
+    from deepmd.pt.utils.finetune import (
         _learned_part_path,
     )
 
diff --git a/source/tests/pt_expt/model/test_get_model_dpa4.py b/source/tests/pt_expt/model/test_get_model_dpa4.py
index 9e2779bacd..ac8d2864a8 100644
--- a/source/tests/pt_expt/model/test_get_model_dpa4.py
+++ b/source/tests/pt_expt/model/test_get_model_dpa4.py
@@ -236,7 +236,6 @@ def test_unsupported_keys_raise(self) -> None:
         and merely warns about.
         """
         cases = {
-            "spin": ({"use_spin": [True, False], "virtual_scale": [0.3]}, "Spin DPA4"),
             "lora": ({"rank": 4}, "`lora` is not supported"),
         }
         for key, (value, msg_regex) in cases.items():
@@ -252,6 +251,18 @@ def test_preset_out_bias_accepted(self) -> None:
         model = get_model(raw)
         self.assertEqual(model.atomic_model.preset_out_bias, {"energy": [None, [1.0]]})
 
+    def test_spin_requires_native_scheme(self) -> None:
+        """DPA4 aliases require an explicit native spin configuration."""
+        for model_type in ("dpa4", "DPA4", "sezm", "SeZM"):
+            for scheme in (None, "invalid"):
+                with self.subTest(model_type=model_type, scheme=scheme):
+                    raw = _make_raw_model_config(type=model_type)
+                    raw["spin"] = {"use_spin": [True, False]}
+                    if scheme is not None:
+                        raw["spin"]["scheme"] = scheme
+                    with self.assertRaisesRegex(NotImplementedError, "require"):
+                        get_model(raw)
+
     def test_native_spin_capability_gate_standard_config(self) -> None:
         """The generic ``supports_native_spin()`` gate rejects a dense descriptor.
 
diff --git a/source/tests/pt_expt/model/test_graph_builder_dispatch.py b/source/tests/pt_expt/model/test_graph_builder_dispatch.py
index a6185aa6f9..8880282820 100644
--- a/source/tests/pt_expt/model/test_graph_builder_dispatch.py
+++ b/source/tests/pt_expt/model/test_graph_builder_dispatch.py
@@ -152,7 +152,7 @@ def test_explicit_nv_rejects_cpu():
         ("cpu", False, False, True, 1, "vesin"),
         ("cpu", False, False, True, 4, "dense"),
         ("cpu", True, False, False, 1, "dense"),
-        ("cuda", True, True, True, 1, "nv"),
+        ("cuda", True, True, True, 1, "vesin"),
         ("cuda", True, True, True, 4, "nv"),
         ("cuda", False, True, True, 1, "vesin"),
         ("cuda", False, True, True, 4, "dense"),
diff --git a/source/tests/pt_expt/model/test_graph_export_with_comm.py b/source/tests/pt_expt/model/test_graph_export_with_comm.py
index aa73a3f46e..92f83c61f4 100644
--- a/source/tests/pt_expt/model/test_graph_export_with_comm.py
+++ b/source/tests/pt_expt/model/test_graph_export_with_comm.py
@@ -430,24 +430,3 @@ def test_native_spin_graph_with_comm_abi() -> None:
     # the spin-specific outputs must survive the with-comm trace
     for key in ("energy", "force", "force_mag", "mask_mag"):
         assert key in keys, f"{key} missing from the with-comm graph outputs"
-
-
-def test_graph_eligibility_guard_is_stated_exactly_once() -> None:
-    """The innermost ``lower_kind="graph"`` guard must not be duplicated.
-
-    It was once present as two verbatim copies from separate commits; the
-    second could never fire yet invited the two to drift. Asserted on the
-    source text -- behavior cannot distinguish one guard from two. Read from
-    the file, not ``inspect.getsource``: conftest wraps ``_trace_and_export``.
-    """
-    from pathlib import (
-        Path,
-    )
-
-    from deepmd.pt_expt.utils import (
-        serialization,
-    )
-
-    source = Path(serialization.__file__).read_text()
-    assert source.count("graph-lower eligible (model_uses_graph_lower() is False") == 1
-    assert source.count("if not model_uses_graph_lower(model):") == 1
diff --git a/source/tests/pt_expt/model/test_legacy_vacuum_ref.py b/source/tests/pt_expt/model/test_legacy_vacuum_ref.py
new file mode 100644
index 0000000000..e6aea5c55b
--- /dev/null
+++ b/source/tests/pt_expt/model/test_legacy_vacuum_ref.py
@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Legacy vacuum references remain differentiable on the graph route."""
+
+import pytest
+import torch
+
+from deepmd.dpmodel.atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.dpmodel.utils.neighbor_graph import (
+    graph_from_dense_quartet,
+)
+from deepmd.pt_expt.common import (
+    auto_wrapped_class,
+)
+from deepmd.pt_expt.utils import (
+    env,
+)
+
+from ...common.dpmodel.test_vacuum_descriptor_graph import (
+    make_model,
+)
+
+
+@pytest.mark.parametrize("kind", ["dpa1", "dpa2"])  # legacy descriptor
+def test_graph_vacuum_reference_preserves_parameter_gradients(kind) -> None:
+    model = auto_wrapped_class(DPAtomicModel).deserialize(make_model(kind).serialize())
+    descriptor = model.descriptor
+    blocks = (
+        [descriptor.se_atten]
+        if kind == "dpa1"
+        else [descriptor.repinit, descriptor.repformers]
+    )
+    with torch.no_grad():
+        for block in blocks:
+            block.mean.fill_(0.3)
+    coord = torch.tensor(
+        [[[0.0, 0.0, 0.0], [1.0, 0.2, 0.0]]],
+        dtype=torch.float64,
+        device=env.DEVICE,
+    )
+    atype = torch.tensor([[0, 1]], dtype=torch.int64, device=env.DEVICE)
+    nlist = torch.tensor([[[1, -1], [0, -1]]], dtype=torch.int64, device=env.DEVICE)
+    graph, atype = graph_from_dense_quartet(coord, atype, nlist, None)
+    gg = descriptor.call_graph(
+        graph, atype, type_embedding=descriptor.graph_type_embedding_table()
+    )[0]
+    reference = model.vacuum_descriptor()
+    parameters = tuple(p for p in model.parameters() if p.requires_grad)
+    reference_grad = torch.autograd.grad(
+        reference.sum(), parameters, allow_unused=True, retain_graph=True
+    )
+    assert any(g is not None and torch.any(g.abs() > 1e-7) for g in reference_grad)
+    expected = model.fitting_net.call_graph(gg, atype, vacuum_descriptor=reference)[
+        "energy"
+    ]
+    actual = model.forward_atomic_graph(graph, atype)["energy"]
+    torch.testing.assert_close(actual, expected, rtol=1e-12, atol=1e-12)
+    expected_grad = torch.autograd.grad(expected.sum(), parameters, allow_unused=True)
+    actual_grad = torch.autograd.grad(actual.sum(), parameters, allow_unused=True)
+    torch.testing.assert_close(actual_grad, expected_grad, rtol=1e-12, atol=1e-12)
diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py
new file mode 100644
index 0000000000..0581dd8626
--- /dev/null
+++ b/source/tests/pt_expt/model/test_long_range.py
@@ -0,0 +1,2798 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Long-range latent-charge modules on the DPA4C graph lower (pt_expt)."""
+
+import copy
+import json
+import os
+import pathlib
+import subprocess
+import sys
+from typing import (
+    Any,
+)
+
+import numpy as np
+import pytest
+import torch
+
+from deepmd.dpmodel.long_range import (
+    COULOMB_CONSTANT,
+)
+from deepmd.dpmodel.utils.neighbor_graph import (
+    ALL_PAIRS,
+    build_pair_graph,
+)
+from deepmd.infer import (
+    DeepPot,
+)
+from deepmd.pt_expt.entrypoints.main import (
+    change_bias,
+)
+from deepmd.pt_expt.model.get_model import (
+    get_model,
+)
+from deepmd.pt_expt.train.training import (
+    _CompiledModel,
+    _get_model_structure_key,
+    _trace_and_compile_graph,
+    model_change_out_bias,
+)
+from deepmd.pt_expt.train.wrapper import (
+    ModelWrapper,
+)
+from deepmd.pt_expt.utils import (
+    env,
+    serialization,
+)
+from deepmd.pt_expt.utils.serialization import (
+    check_mesh_compile_torch_version,
+    check_mesh_freeze_torch_version,
+    deserialize_to_file,
+    mesh_compile_supported,
+    mesh_freeze_supported,
+    serialize_from_file,
+)
+
+from ..compile_utils import (
+    REQUIRES_SUPPORTED_COMPILE,
+)
+
+KERNELS = {
+    "les": {"type": "les", "alpha": 0.8},
+    "sog": {"type": "sog", "amplitude": [0.7, -0.2, 0.1], "width": [0.8, 1.7, 3.2]},
+}
+CHANNELS = {"les": 1, "sog": 3}
+NEUTRAL = {"source": "neutral"}
+
+
+def _config(kernel: str = "les", constraint: dict | None = NEUTRAL) -> dict:
+    """Small float64 DPA4C model with a latent-charge long-range module."""
+    return {
+        "type_map": ["A", "B"],
+        "descriptor": {
+            "type": "dpa4c",
+            "rcut": 3.0,
+            "channels": 8,
+            "lmax": 2,
+            "n_radial": 4,
+            "precision": "float64",
+            "seed": 17,
+        },
+        "fitting_net": {
+            "type": "ener",
+            "neuron": [16, 16],
+            "precision": "float64",
+            "seed": 19,
+        },
+        "long_range": {
+            "type": "latent_charge",
+            "channels": CHANNELS[kernel],
+            "charge_net": {"neuron": [8, 8], "precision": "float64", "seed": 23},
+            "kernel": copy.deepcopy(KERNELS[kernel]),
+            "charge_constraint": copy.deepcopy(constraint),
+        },
+    }
+
+
+def _model(kernel: str = "les", **kwargs: object) -> torch.nn.Module:
+    return get_model(_config(kernel, **kwargs)).to(env.DEVICE).eval()
+
+
+def _tensor(x: np.ndarray, dtype: torch.dtype = torch.float64) -> torch.Tensor:
+    return torch.tensor(x, dtype=dtype, device=env.DEVICE)
+
+
+def _batch(seed: int = 0) -> tuple[torch.Tensor, torch.Tensor]:
+    """Three frames of 6, 3 and 5 real atoms, padded to a common width."""
+    rng = np.random.default_rng(seed)
+    coord = rng.uniform(0.0, 6.0, (3, 6, 3))
+    atype = np.array(
+        [[0, 1, 0, 1, 1, 0], [1, 0, 0, -1, -1, -1], [0, 0, 1, 1, 0, -1]],
+        dtype=np.int64,
+    )
+    return _tensor(coord), _tensor(atype, torch.int64)
+
+
+def _set_constant_charges(model: torch.nn.Module, charges: list[float]) -> None:
+    """Zero every network weight and give the atoms of each type a fixed charge.
+
+    The energy is then the long-range pair sum alone, of charges that do not
+    depend on the geometry; the kernel keeps its configured parameters.
+    """
+    with torch.no_grad():
+        for name, parameter in model.atomic_model.named_parameters():
+            if ".kernel." not in f".{name}":
+                parameter.zero_()
+        bias = model.atomic_model.long_range.charge_net.bias_atom_e
+        bias.zero_()
+        bias[:, 0] = torch.as_tensor(charges, dtype=bias.dtype)
+
+
+def _reference_energy(coord: np.ndarray, charge: np.ndarray, kernel: str) -> float:
+    """Direct pair sum of the latent-charge energy of one frame, in eV."""
+    energy = 0.0
+    for i in range(len(coord)):
+        for j in range(i + 1, len(coord)):
+            r = np.linalg.norm(coord[i] - coord[j])
+            if kernel == "les":
+                from scipy.special import (
+                    erf,
+                )
+
+                k = erf(0.8 * r) / r
+            else:
+                amplitude = np.array([0.7, -0.2, 0.1])
+                width = np.array([0.8, 1.7, 3.2])
+                k = np.sum(amplitude * np.exp(-(r**2) / (2 * width**2)))
+            energy += COULOMB_CONSTANT * float(charge[i] @ charge[j]) * k
+    return energy
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_energy_matches_direct_pair_sum(kernel: str) -> None:
+    """With fixed charges the energy is the documented pair sum, attributed half per atom."""
+    model = _model(kernel, constraint=None)
+    _set_constant_charges(model, [1.0, -0.5])
+    coord = np.array([[[0.0, 0.0, 0.0], [1.3, 0.0, 0.0], [0.0, 2.1, 0.4]]])
+    atype = np.array([[0, 1, 1]])
+    out = model(_tensor(coord), _tensor(atype, torch.int64))
+    charge = np.zeros((3, CHANNELS[kernel]))
+    charge[:, 0] = [1.0, -0.5, -0.5]
+    expected = _reference_energy(coord[0], charge, kernel)
+    np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-12)
+    np.testing.assert_allclose(
+        out["atom_energy"].sum().item(), expected, rtol=1e-12, atol=1e-14
+    )
+
+
+@pytest.mark.parametrize("condition", ["aparam", "case_embd"])
+def test_charges_follow_the_fitting_conditions(condition: str) -> None:
+    """The charge network reads the atomic parameters and the case of the energy fitting.
+
+    The long-range energy is the model less the short-range model of the same
+    seeds; it changes with the atomic parameters of the atoms and with the
+    case that a multi-task model selects.
+    """
+    config = _config("les", constraint=None)
+    # Two atomic parameters, or the one-hot embedding of two cases.
+    key = "numb_aparam" if condition == "aparam" else "dim_case_embd"
+    config["fitting_net"][key] = 2
+    short_range_config = copy.deepcopy(config)
+    short_range_config.pop("long_range")
+    model = get_model(config).to(env.DEVICE).eval()
+    # The same seeds give the short-range networks of both models.
+    short_range = get_model(short_range_config).to(env.DEVICE).eval()
+    neutral = copy.deepcopy(model)
+    with torch.no_grad():
+        for parameter in neutral.atomic_model.long_range.charge_net.parameters():
+            parameter.zero_()
+    coord, atype = _batch(seed=7)
+    long_range_energy = []
+    for value in (0, 1):
+        kwargs = {}
+        if condition == "aparam":
+            kwargs["aparam"] = torch.full(
+                (*atype.shape, 2), 1.5 * value, dtype=torch.float64, device=env.DEVICE
+            )
+        else:
+            for each in (model, short_range, neutral):
+                each.set_case_embd(value)
+        energy = short_range(coord, atype, **kwargs)["energy"]
+        # Without charges the long-range model is the short-range one.
+        torch.testing.assert_close(neutral(coord, atype, **kwargs)["energy"], energy)
+        long_range_energy.append(model(coord, atype, **kwargs)["energy"] - energy)
+    difference = (long_range_energy[1] - long_range_energy[0]).abs()
+    assert torch.all(difference > 1e-6), difference
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_force_virial_hessian_match_finite_differences(kernel: str) -> None:
+    """Force, affine-strain virial and Hessian are derivatives of the complete energy."""
+    model = _model(kernel)
+    # Atoms within the cutoff of each other make the charges depend on the
+    # geometry, so the derivatives hold the response of the charges.
+    rng = np.random.default_rng(3)
+    coord = rng.uniform(0.0, 2.5, (1, 4, 3))
+    atype = np.array([[0, 1, 0, 1]])
+    coord_t, atype_t = _tensor(coord), _tensor(atype, torch.int64)
+    out = model(coord_t, atype_t)
+
+    def energy(c: np.ndarray) -> float:
+        return model(_tensor(c), atype_t)["energy"].item()
+
+    h = 1e-5
+    force_fd = np.zeros((4, 3))
+    for i in range(4):
+        for d in range(3):
+            plus, minus = coord.copy(), coord.copy()
+            plus[0, i, d] += h
+            minus[0, i, d] -= h
+            force_fd[i, d] = -(energy(plus) - energy(minus)) / (2 * h)
+    np.testing.assert_allclose(
+        out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-8
+    )
+
+    # The virial is -dE/d(strain) at zero strain, for x -> (1 + eps) x.
+    virial_fd = np.zeros((3, 3))
+    for a in range(3):
+        for b in range(3):
+            strain = np.zeros((3, 3))
+            strain[a, b] = h
+            plus = coord @ (np.eye(3) + strain)
+            minus = coord @ (np.eye(3) - strain)
+            virial_fd[a, b] = -(energy(plus) - energy(minus)) / (2 * h)
+    np.testing.assert_allclose(
+        out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-8
+    )
+
+    hessian_model = _model(kernel)
+    hessian_model.enable_hessian()
+    hessian = hessian_model(coord_t, atype_t)["hessian"].detach().cpu().numpy()
+    hessian = hessian.reshape(12, 12)
+    hessian_fd = np.zeros((12, 12))
+    for k in range(12):
+        plus, minus = coord.copy(), coord.copy()
+        plus.reshape(-1)[k] += h
+        minus.reshape(-1)[k] -= h
+        f_plus = model(_tensor(plus), atype_t)["force"].detach().cpu().numpy()
+        f_minus = model(_tensor(minus), atype_t)["force"].detach().cpu().numpy()
+        hessian_fd[:, k] = -(f_plus - f_minus).reshape(-1) / (2 * h)
+    np.testing.assert_allclose(hessian, hessian_fd, atol=1e-6)
+    np.testing.assert_allclose(hessian, hessian.T, atol=1e-10)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_atomic_reductions_and_padding(kernel: str) -> None:
+    """Atomic energy and virial sum to the frame values; padding atoms are inert."""
+    model = _model(kernel)
+    coord, atype = _batch()
+    out = model(coord, atype, do_atomic_virial=True)
+    torch.testing.assert_close(out["atom_energy"].sum(1), out["energy"])
+    torch.testing.assert_close(out["atom_virial"].sum(1), out["virial"])
+    assert out["force"][1, 3:].abs().max() == 0
+    moved = coord.clone()
+    moved[1, 3:] += 50.0
+    moved_out = model(moved, atype)
+    torch.testing.assert_close(moved_out["energy"], out["energy"], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_frames_are_independent(kernel: str) -> None:
+    """A frame keeps its energy and force whatever other frames are batched with it."""
+    model = _model(kernel)
+    coord, atype = _batch()
+    batched = model(coord, atype)
+    for frame, nloc in enumerate((6, 3, 5)):
+        alone = model(coord[frame : frame + 1, :nloc], atype[frame : frame + 1, :nloc])
+        torch.testing.assert_close(alone["energy"][0], batched["energy"][frame])
+        torch.testing.assert_close(
+            alone["force"][0], batched["force"][frame, :nloc], atol=1e-12, rtol=0
+        )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_ragged_matches_rectangular(kernel: str) -> None:
+    """The flat ragged entry evaluates the same model as the padded one."""
+    model = _model(kernel)
+    coord, atype = _batch()
+    padded = model(coord, atype)
+    keep = atype.reshape(-1) >= 0
+    ragged = model.forward_ragged(
+        coord.reshape(-1, 3)[keep],
+        atype.reshape(-1)[keep],
+        torch.tensor([6, 3, 5], device=env.DEVICE),
+    )
+    torch.testing.assert_close(ragged["energy"], padded["energy"])
+    torch.testing.assert_close(ragged["force"], padded["force"].reshape(-1, 3)[keep])
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_empty_and_single_atom_frames(kernel: str) -> None:
+    """Frames without pairs contribute no long-range energy and stay finite."""
+    model = _model(kernel)
+    coord = _tensor(np.random.default_rng(5).uniform(0.0, 3.0, (3, 2, 3)))
+    atype = _tensor(np.array([[0, -1], [1, 0], [-1, -1]]), torch.int64)
+    out = model(coord, atype)
+    assert torch.isfinite(out["energy"]).all()
+    assert torch.isfinite(out["force"]).all()
+    single = model(coord[:1, :1], atype[:1, :1])
+    torch.testing.assert_close(single["energy"][0], out["energy"][0])
+
+
+def test_permutation_and_rigid_motion_invariance() -> None:
+    """The energy is invariant under atom permutation, translation and rotation."""
+    model = _model("sog")
+    coord, atype = _batch()
+    coord, atype = coord[:1], atype[:1]
+    out = model(coord, atype)
+    perm = torch.tensor([3, 0, 5, 1, 4, 2], device=env.DEVICE)
+    permuted = model(coord[:, perm], atype[:, perm])
+    torch.testing.assert_close(permuted["energy"], out["energy"])
+    torch.testing.assert_close(permuted["force"], out["force"][:, perm])
+    rotation, _ = torch.linalg.qr(
+        _tensor(np.random.default_rng(13).normal(size=(3, 3)))
+    )
+    moved = model(coord @ rotation.T + 3.7, atype)
+    torch.testing.assert_close(moved["energy"], out["energy"])
+    torch.testing.assert_close(moved["force"], out["force"] @ rotation.T)
+
+
+@pytest.mark.parametrize("source", ["charge_spin", "fparam"])
+def test_constraint_default_matches_explicit_target(source: str) -> None:
+    """An omitted frame condition takes its default for the constraint as well."""
+    config = _config("les", constraint={"source": source, "fparam_index": 0})
+    if source == "charge_spin":
+        config["descriptor"]["add_chg_spin_ebd"] = True
+        config["descriptor"]["default_chg_spin"] = [2.0, 1.0]
+        explicit = {"charge_spin": _tensor(np.array([[2.0, 1.0], [2.0, 1.0]]))}
+        other = {"charge_spin": _tensor(np.array([[0.0, 1.0], [0.0, 1.0]]))}
+    else:
+        config["fitting_net"]["numb_fparam"] = 1
+        config["fitting_net"]["default_fparam"] = [2.0]
+        explicit = {"fparam": _tensor(np.array([[2.0], [2.0]]))}
+        other = {"fparam": _tensor(np.array([[0.0], [0.0]]))}
+    model = get_model(config).to(env.DEVICE).eval()
+    coord, atype = _batch()
+    coord, atype = coord[:2], atype[:2]
+    implicit_out = model(coord, atype)
+    explicit_out = model(coord, atype, **explicit)
+    torch.testing.assert_close(implicit_out["energy"], explicit_out["energy"])
+    torch.testing.assert_close(implicit_out["force"], explicit_out["force"])
+    other_out = model(coord, atype, **other)
+    assert (other_out["energy"] - implicit_out["energy"]).abs().max() > 1e-6
+
+
+@pytest.mark.parametrize("constraint", [None, NEUTRAL])
+def test_pair_exclusion_removes_the_interaction(constraint: dict | None) -> None:
+    """Pairs of excluded types do not interact; their atoms keep their charges."""
+    config = _config("les", constraint=constraint)
+    config["pair_exclude_types"] = [[0, 1]]
+    model = get_model(config).to(env.DEVICE).eval()
+    _set_constant_charges(model, [1.0, -1.0])
+    coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]]))
+    atype = _tensor(np.array([[0, 0, 1]]), torch.int64)
+    out = model(coord, atype)
+    # Only the two type-0 atoms interact. The neutral constraint shifts the
+    # charges (1, 1, -1) of the three atoms to (2/3, 2/3, -4/3).
+    charge = 1.0 if constraint is None else 2.0 / 3.0
+    erf = torch.erf(torch.tensor(3.2, dtype=torch.float64)).item()
+    expected = COULOMB_CONSTANT * charge**2 * erf / 4.0
+    np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-12)
+    assert out["force"][0, 2].abs().max() == 0
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_les_kernel_derivatives_down_to_contact(dtype: torch.dtype) -> None:
+    """The radial derivatives of LES keep the precision of the dtype near contact.
+
+    The derivatives of erf(alpha r) / r subtract terms that agree to second
+    order in alpha r, which the series of the kernel avoids; the reference
+    differentiates in 40-digit arithmetic.
+    """
+    import mpmath
+
+    from deepmd.pt_expt.long_range.kernel import (
+        LESKernel,
+    )
+
+    alpha = 0.8
+    kernel = LESKernel(alpha=alpha)
+    r = torch.tensor(
+        np.geomspace(1e-5, 6.0, 80), dtype=dtype, device=env.DEVICE
+    ).requires_grad_(True)
+    (first,) = torch.autograd.grad(kernel(r * r).sum(), r, create_graph=True)
+    (second,) = torch.autograd.grad(first.sum(), r)
+
+    def les(t: mpmath.mpf) -> mpmath.mpf:
+        return mpmath.erf(alpha * t) / t
+
+    rtol = 2e-5 if dtype == torch.float32 else 1e-11
+    for order, actual in ((1, first), (2, second)):
+        with mpmath.workdps(40):
+            expected = np.array(
+                [
+                    float(mpmath.diff(les, mpmath.mpf(float(t)), order))
+                    for t in r.detach().cpu()
+                ]
+            )
+        # The second derivative changes sign near 1.2 Å, where it is the
+        # difference of terms much larger than itself.
+        np.testing.assert_allclose(
+            actual.detach().cpu().numpy(),
+            expected,
+            rtol=rtol,
+            atol=rtol * np.abs(expected).max(),
+            err_msg=str(order),
+        )
+
+
+def test_owned_frames_across_domains() -> None:
+    """Halo nodes are rejected on a single domain and kept across domains."""
+    from deepmd.pt_expt.model.make_model import (
+        check_owned_frames,
+    )
+
+    n_node = torch.tensor([5, 3], dtype=torch.int64)
+    n_local = torch.tensor([4, 3], dtype=torch.int64)
+    check_owned_frames(n_node, n_node.clone())
+    check_owned_frames(n_node, n_local, torch.ones(1, dtype=torch.int64))
+    for communicator in (None, torch.zeros(1, dtype=torch.int64)):
+        with pytest.raises(ValueError, match="n_local == n_node"):
+            check_owned_frames(n_node, n_local, communicator)
+
+
+def test_graph_lower_rejects_halo_nodes() -> None:
+    """A graph whose frames carry halo nodes is rejected, not summed without them.
+
+    The halo nodes of a graph without communication would carry no charge and
+    remove their interactions from the sum of a frame held whole.
+    """
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        build_topologies,
+    )
+    from deepmd.pt_expt.utils.graph_builder import (
+        build_neighbor_graph_for_method,
+    )
+
+    model = _model("les", constraint=None)
+    coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]]))
+    atype = _tensor(np.array([[0, 0, 1]]), torch.int64)
+    graph = build_neighbor_graph_for_method(
+        "dense", coord, atype, None, model.get_rcut(), None
+    )
+    topologies = build_topologies(
+        model.atomic_model.required_topologies(),
+        coord.reshape(-1, 3),
+        graph.n_node,
+        atype=atype.reshape(-1),
+    )
+    inputs = (
+        atype.reshape(-1),
+        graph.n_node,
+        graph.edge_index,
+        graph.edge_vec,
+        graph.edge_mask,
+    )
+
+    def lower(n_local: torch.Tensor) -> dict[str, torch.Tensor]:
+        return model.forward_common_lower_graph(
+            inputs[0], inputs[1], n_local, *inputs[2:], topologies=topologies
+        )
+
+    lower(graph.n_node.clone())
+    with pytest.raises(ValueError, match="n_local == n_node"):
+        lower(graph.n_node - 1)
+
+
+def test_linear_composition_takes_the_exclusions_of_the_composition() -> None:
+    """A long-range child follows the pair exclusion of its composition.
+
+    The children of a composition share its neighbor graph and pair topology,
+    so the exclusions of the composition give the child the semantics of the
+    same model alone, and an exclusion of the child itself is rejected.
+    """
+    child = _config("les", constraint=None)
+    child["type"] = "standard"
+    composition = {
+        "type": "linear_ener",
+        "type_map": child["type_map"],
+        "weights": "sum",
+        "models": [child],
+        "pair_exclude_types": [[0, 1]],
+    }
+    model = get_model(copy.deepcopy(composition)).to(env.DEVICE).eval()
+    with torch.no_grad():
+        for name, parameter in model.atomic_model.named_parameters():
+            if ".kernel." not in f".{name}":
+                parameter.zero_()
+        bias = model.atomic_model.models[0].long_range.charge_net.bias_atom_e
+        bias[:, 0] = torch.as_tensor([1.0, -1.0], dtype=bias.dtype)
+    coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]]))
+    atype = _tensor(np.array([[0, 0, 1]]), torch.int64)
+    # As in test_pair_exclusion_removes_the_interaction: only the two type-0
+    # atoms interact.
+    erf = torch.erf(torch.tensor(3.2, dtype=torch.float64)).item()
+    np.testing.assert_allclose(
+        model(coord, atype)["energy"].item(),
+        COULOMB_CONSTANT * erf / 4.0,
+        rtol=1e-12,
+    )
+    for key, value, error, match in (
+        ("pair_exclude_types", [[0, 1]], ValueError, "from the composition"),
+        ("atom_exclude_types", [1], NotImplementedError, "atom_exclude_types"),
+    ):
+        config = copy.deepcopy(composition)
+        config.pop("pair_exclude_types")
+        config["models"][0][key] = value
+        with pytest.raises(error, match=match):
+            get_model(config)
+    # a periodic module shares its composition with no other long-range module
+    periodic = copy.deepcopy(child)
+    periodic["long_range"]["periodic"] = True
+    for models in ([periodic, child], [periodic, periodic]):
+        config = copy.deepcopy(composition)
+        config.pop("pair_exclude_types")
+        config["models"] = copy.deepcopy(models)
+        with pytest.raises(NotImplementedError, match="serves periodic systems"):
+            get_model(config)
+    config["models"] = [copy.deepcopy(child), copy.deepcopy(child)]
+    get_model(config)
+    # a periodic module takes no pair exclusion, alone or in a composition
+    config = copy.deepcopy(composition)
+    config["models"] = [copy.deepcopy(periodic)]
+    with pytest.raises(NotImplementedError, match="pair_exclude_types"):
+        get_model(config)
+    config = copy.deepcopy(periodic)
+    config["type"] = "standard"
+    config["pair_exclude_types"] = [[0, 1]]
+    with pytest.raises(NotImplementedError, match="pair_exclude_types"):
+        get_model(config)
+
+
+@pytest.mark.parametrize("compiled", [False, True])
+def test_vacuum_reference_with_long_range(compiled: bool) -> None:
+    """An isolated atom keeps its reference energy; reference rows join no pair."""
+    config = _config("les", constraint={"source": "charge_spin"})
+    config["type_map"] = ["O", "H"]
+    config["preset_out_bias"] = {"energy": {"O": -3.0, "H": 0.5}}
+    config["descriptor"]["add_chg_spin_ebd"] = True
+    config["descriptor"]["default_chg_spin"] = [0.0, 1.0]
+    config["fitting_net"]["vacuum_ref"] = True
+    model = get_model(config).to(env.DEVICE).eval()
+    assert model.get_fitting_net().needs_vacuum_descriptor()
+    rng = np.random.default_rng(11)
+    coord = _tensor(rng.uniform(0.0, 2.5, (2, 4, 3)))
+    atype = _tensor(np.array([[0, 1, 1, 0], [1, 0, -1, -1]]), torch.int64)
+    if compiled:
+        model.train()
+        runner = _CompiledModel(model, _get_model_structure_key(model))
+        runner.train()
+        reference = model(coord, atype)
+        out = runner(coord, atype)
+        torch.testing.assert_close(out["energy"], reference["energy"])
+        torch.testing.assert_close(out["force"], reference["force"])
+        return
+    out = model(coord, atype)
+    for frame in range(2):
+        alone = model(coord[frame : frame + 1], atype[frame : frame + 1])
+        torch.testing.assert_close(alone["energy"][0], out["energy"][frame])
+    # The output statistics set the bias; an atom without neighbors and in
+    # its reference state then has exactly the bias of its type.
+    bias = model.atomic_model.out_bias.clone()
+    bias[0, :, 0] = torch.tensor([-3.0, 0.5])
+    model.atomic_model.set_out_bias(bias)
+    isolated = model(
+        coord[:1, :1],
+        atype[:1, :1],
+        charge_spin=_tensor(np.array([[0.0, 3.0]])),
+    )
+    torch.testing.assert_close(
+        isolated["energy"].reshape(()),
+        torch.tensor(-3.0, dtype=torch.float64, device=env.DEVICE),
+        rtol=0,
+        atol=1e-12,
+    )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_self_label_bias_adjustment_is_invariant(kernel: str) -> None:
+    """Labelling a sample with the model's own energies leaves its bias unchanged."""
+    model = _model(kernel)
+    with torch.no_grad():
+        model.atomic_model.long_range.charge_net.bias_atom_e[:, 0] = torch.tensor(
+            [0.8, -0.6]
+        )
+    coord, atype = _batch()
+    coord, atype = coord[:1], atype[:1]
+    energy = model(coord, atype)["energy"].detach().cpu().numpy()
+    natoms = np.array([[6, 6, 3, 3]], dtype=np.int32)
+    sample = [
+        {
+            "coord": coord.cpu().numpy(),
+            "atype": atype.cpu().numpy(),
+            "box": None,
+            "natoms": natoms,
+            "energy": energy.reshape(1, 1),
+            "find_energy": np.float32(1.0),
+        }
+    ]
+    bias_before = model.atomic_model.out_bias.detach().clone()
+    model_change_out_bias(model, lambda: sample, "change-by-statistic")
+    torch.testing.assert_close(
+        model.atomic_model.out_bias, bias_before, atol=1e-10, rtol=0
+    )
+    after = model(coord, atype)["energy"].detach().cpu().numpy()
+    np.testing.assert_allclose(after, energy, atol=1e-10)
+
+
+@pytest.mark.parametrize("weights", ["sum", "mean", [0.3, 1.4]])
+def test_linear_composition_combines_complete_energies(weights: object) -> None:
+    """A linear composition weights the complete energy of each child.
+
+    The second child is an ordinary model, which reads no topology.
+    """
+    child = _config("les")
+    child["type"] = "standard"
+    other = copy.deepcopy(child)
+    other.pop("long_range")
+    other["descriptor"]["seed"] = 31
+    composition = (
+        get_model(
+            {
+                "type": "linear_ener",
+                "type_map": child["type_map"],
+                "weights": weights,
+                "models": [child, other],
+            }
+        )
+        .to(env.DEVICE)
+        .eval()
+    )
+    standalone = [_model("les"), get_model(other).to(env.DEVICE).eval()]
+    for index, model in enumerate(standalone):
+        model.atomic_model.load_state_dict(
+            composition.atomic_model.models[index].state_dict()
+        )
+    coord, atype = _batch()
+    out = composition(coord, atype)
+    if weights == "sum":
+        w = [1.0, 1.0]
+    elif weights == "mean":
+        w = [0.5, 0.5]
+    else:
+        w = weights
+    energy = sum(
+        ww * m(coord, atype)["energy"] for ww, m in zip(w, standalone, strict=True)
+    )
+    force = sum(
+        ww * m(coord, atype)["force"] for ww, m in zip(w, standalone, strict=True)
+    )
+    torch.testing.assert_close(out["energy"], energy)
+    torch.testing.assert_close(out["force"], force)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_parameter_gradients_and_freezing(kernel: str) -> None:
+    """Energy and force losses reach the charge network and the kernel; frozen parts stay out.
+
+    Every parameter of the module moves the losses, except under a constraint
+    the common charge bias of the constrained channel, whose gradient
+    vanishes (see test_constraint_removes_the_common_charge_bias).
+    """
+    coord, atype = _batch()
+    for constraint in (None, NEUTRAL):
+        model = _model(kernel, constraint=constraint).train()
+        out = model(coord, atype)
+        loss = out["energy"].square().sum() + out["force"].square().sum()
+        loss.backward()
+        long_range = model.atomic_model.long_range
+        bias = long_range.charge_net.nets[0].layers[-1].b
+        for name, parameter in long_range.named_parameters():
+            assert parameter.grad is not None, name
+            assert torch.isfinite(parameter.grad).all(), name
+            gradient = parameter.grad
+            if constraint is not None and parameter is bias:
+                # Channel 0 is the constrained channel.
+                gradient = gradient[..., 1:]
+            if gradient.numel():
+                assert gradient.abs().max() > 0, name
+
+    config = _config(kernel)
+    config["long_range"]["kernel"]["trainable"] = False
+    config["long_range"]["charge_net"]["trainable"] = False
+    frozen = get_model(config).to(env.DEVICE).train()
+    names = [
+        name
+        for name, parameter in frozen.named_parameters()
+        if parameter.requires_grad and "long_range" in name
+    ]
+    assert names == []
+    frozen(coord, atype)["energy"].sum().backward()
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_constraint_removes_the_common_charge_bias(kernel: str) -> None:
+    """A uniform shift of the raw charges of the constrained channel changes nothing.
+
+    The constraint subtracts the mean shift of every frame, so the output bias
+    of that channel leaves energy and force unchanged, while that of another
+    channel does not.
+    """
+    model = _model(kernel)
+    coord, atype = _batch()
+    reference = model(coord, atype)
+    bias = model.atomic_model.long_range.charge_net.nets[0].layers[-1].b
+    for channel in range(bias.shape[-1]):
+        with torch.no_grad():
+            bias[..., channel] += 0.3
+        out = model(coord, atype)
+        with torch.no_grad():
+            bias[..., channel] -= 0.3
+        if channel == 0:
+            torch.testing.assert_close(out["energy"], reference["energy"])
+            torch.testing.assert_close(out["force"], reference["force"])
+        else:
+            assert (out["energy"] - reference["energy"]).abs().max() > 1e-6
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_serialization_round_trip(kernel: str) -> None:
+    """The portable record and the state dict restore the same model."""
+    model = _model(kernel)
+    rng = np.random.default_rng(5)
+    with torch.no_grad():
+        for parameter in model.atomic_model.long_range.parameters():
+            parameter.add_(0.01 * _tensor(rng.normal(size=tuple(parameter.shape))))
+    coord, atype = _batch()
+    expected = model(coord, atype)
+    record = model.serialize()
+    assert record["@version"] == 3
+    restored = type(model).deserialize(copy.deepcopy(record)).to(env.DEVICE).eval()
+    torch.testing.assert_close(restored(coord, atype)["energy"], expected["energy"])
+    rebuilt = _model(kernel)
+    rebuilt.load_state_dict(model.state_dict())
+    torch.testing.assert_close(rebuilt(coord, atype)["force"], expected["force"])
+
+
+def test_construction_is_reproducible() -> None:
+    """Without a seed of its own the charge network follows the fitting seed."""
+    config = _config("les")
+    del config["long_range"]["charge_net"]["seed"]
+    first, second = (get_model(copy.deepcopy(config)) for _ in range(2))
+    for (name, a), (_, b) in zip(
+        first.named_parameters(), second.named_parameters(), strict=True
+    ):
+        torch.testing.assert_close(a, b, rtol=0, atol=0, msg=name)
+
+
+def test_ordinary_record_keeps_version() -> None:
+    """A model without a long-range module keeps its record and constructor."""
+    config = _config("les")
+    config.pop("long_range")
+    model = get_model(config)
+    record = model.serialize()
+    assert record["@version"] == 2
+    assert "long_range" not in record
+    assert not model.atomic_model.has_long_range()
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_checkpoint_inference_matches_direct_forward(tmp_path, kernel: str) -> None:
+    """DeepPot on a checkpoint builds the topology over the real atoms only."""
+    model = _model(kernel)
+    checkpoint = tmp_path / f"long_range_{kernel}.pt"
+    wrapper = ModelWrapper(model, model_params=_config(kernel))
+    torch.save({"model": wrapper.state_dict()}, checkpoint)
+    coord, atype = _batch(seed=7)
+    expected = model(coord, atype, do_atomic_virial=True)
+    dp = DeepPot(str(checkpoint), auto_batch_size=False)
+    names = ("energy", "force", "virial", "atom_energy", "atom_virial")
+    coord_np = coord.cpu().numpy()
+    actual = dict(
+        zip(
+            names,
+            dp.eval(coord_np, None, atype.cpu().numpy(), atomic=True, mixed_type=True),
+            strict=True,
+        )
+    )
+    for name in names:
+        np.testing.assert_allclose(
+            actual[name],
+            expected[name].detach().cpu().numpy(),
+            rtol=1e-10,
+            atol=1e-10,
+            err_msg=name,
+        )
+    coord_np[1, 3:] += 40.0
+    moved = dp.eval(coord_np, None, atype.cpu().numpy(), mixed_type=True)
+    np.testing.assert_allclose(moved[0], actual["energy"], rtol=0, atol=1e-12)
+
+
+@pytest.mark.parametrize("grid_mode", ["ragged", "pme"])
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_checkpoint_inference_matches_direct_forward(
+    tmp_path, kernel: str, grid_mode: str
+) -> None:
+    """DeepPot on a periodic checkpoint builds the periodic topology and matches."""
+    if grid_mode == "pme":
+        model = _periodic_pme_model(kernel)
+        config = _config(kernel)
+        config["long_range"]["periodic"] = True
+        config["long_range"]["k_grid_mode"] = "pme"
+    else:
+        model = _periodic_model(kernel)
+        config = _config(kernel)
+        config["long_range"]["periodic"] = True
+        config["long_range"]["k_grid_mode"] = "ragged"
+    checkpoint = tmp_path / f"periodic_{kernel}_{grid_mode}.pt"
+    wrapper = ModelWrapper(model, model_params=config)
+    torch.save({"model": wrapper.state_dict()}, checkpoint)
+    coord, atype, box = _periodic_case()
+    coord_t, atype_t, box_t = (
+        _tensor(coord),
+        _tensor(atype, torch.int64),
+        _tensor(box),
+    )
+    expected = model(coord_t, atype_t, box=box_t, do_atomic_virial=True)
+    dp = DeepPot(str(checkpoint), auto_batch_size=False)
+    names = ("energy", "force", "virial", "atom_energy", "atom_virial")
+    actual = dict(
+        zip(
+            names,
+            dp.eval(
+                coord,
+                box,
+                atype,
+                atomic=True,
+                mixed_type=True,
+            ),
+            strict=True,
+        )
+    )
+    for name in names:
+        np.testing.assert_allclose(
+            actual[name],
+            expected[name].detach().cpu().numpy(),
+            rtol=1e-9,
+            atol=1e-9,
+            err_msg=name,
+        )
+
+
+@pytest.mark.parametrize(
+    ("kernel", "source", "suffix", "atomic"),
+    [
+        ("les", "neutral", ".pt2", False),
+        ("sog", "neutral", ".pte", False),
+        ("les", "fparam", ".pte", False),
+        ("les", "neutral", ".pte", True),
+    ],
+)
+def test_frozen_model_matches_direct_forward(
+    tmp_path, kernel: str, source: str, suffix: str, atomic: bool
+) -> None:
+    """A frozen model reads the pair topology as named inputs of any length.
+
+    With ``atomic``, the model reads atomic parameters, which reach the charge
+    network, and excludes the pairs of two types, which the deployed pair
+    topology omits.
+    """
+    config = _config(kernel, constraint={"source": source, "fparam_index": 0})
+    frame = {}
+    if source == "fparam":
+        config["fitting_net"]["numb_fparam"] = 1
+        frame["fparam"] = np.array([[1.0], [-0.5], [2.0]])
+    if atomic:
+        config["fitting_net"]["numb_aparam"] = 1
+        config["pair_exclude_types"] = [[0, 1]]
+    model = get_model(config).to(env.DEVICE).eval()
+    frozen = tmp_path / f"long_range_{kernel}{suffix}"
+    deserialize_to_file(
+        str(frozen),
+        {"model": model.serialize(), "model_def_script": config},
+        do_atomic_virial=True,
+        lower_kind="graph",
+    )
+    dp = DeepPot(str(frozen), auto_batch_size=False)
+    assert dp.deep_eval.metadata["topologies"] == ["all_pairs"]
+    # The archive reports its lower and its per-atom virial, which
+    # change-bias carries over.
+    recorded = serialize_from_file(str(frozen))
+    assert recorded["lower_input_kind"] == "graph"
+    assert recorded["do_atomic_virial"]
+    if suffix == ".pte":
+        # The pair topology follows the thirteen inputs of the graph lower.
+        signature = torch.export.load(str(frozen)).graph_signature
+        assert len(signature.user_inputs) == 13 + 3
+    names = ("energy", "force", "virial", "atom_energy", "atom_virial")
+    for seed, nloc in ((7, 6), (8, 4)):
+        coord, atype = _batch(seed=seed)
+        coord, atype = coord[:, :nloc], atype[:, :nloc]
+        if atomic:
+            rng = np.random.default_rng(seed)
+            frame["aparam"] = rng.uniform(-1.0, 1.0, (*atype.shape, 1))
+        expected = model(
+            coord,
+            atype,
+            do_atomic_virial=True,
+            **{key: _tensor(value) for key, value in frame.items()},
+        )
+        actual = dp.eval(
+            coord.cpu().numpy(),
+            None,
+            atype.cpu().numpy(),
+            atomic=True,
+            mixed_type=True,
+            **frame,
+        )
+        for name, value in zip(names, actual, strict=True):
+            np.testing.assert_allclose(
+                value,
+                expected[name].detach().cpu().numpy().reshape(value.shape),
+                rtol=1e-10,
+                atol=1e-10,
+                err_msg=name,
+            )
+
+
+@pytest.mark.parametrize(
+    ("suffix", "recorded"), [(".pte", "graph"), (".pt2", "dpa4c_canonical")]
+)
+def test_change_bias_keeps_the_graph_lower(
+    tmp_path, monkeypatch, suffix: str, recorded: str
+) -> None:
+    """change-bias rewrites a frozen long-range model on its graph lower.
+
+    A .pt2 archive that records the compact canonical lower, a form of CUDA
+    devices, is rewritten on the lower that the host selects, the graph lower
+    for this uncompressed model on any device.
+    """
+    config = _config("les")
+    model = get_model(config).to(env.DEVICE).eval()
+    frozen = tmp_path / f"long_range{suffix}"
+    deserialize_to_file(
+        str(frozen),
+        {"model": model.serialize(), "model_def_script": config},
+        lower_kind="graph",
+    )
+    if recorded != "graph":
+        read = serialization.serialize_from_file
+        monkeypatch.setattr(
+            serialization,
+            "serialize_from_file",
+            lambda path: {**read(path), "lower_input_kind": recorded},
+        )
+    bias = [0.5, -1.25]
+    updated = tmp_path / f"long_range_updated{suffix}"
+    change_bias(str(frozen), bias_value=bias, output=str(updated))
+    before = DeepPot(str(frozen), auto_batch_size=False)
+    after = DeepPot(str(updated), auto_batch_size=False)
+    assert after.deep_eval.metadata["lower_input_kind"] == "graph"
+    coord, atype = _batch(seed=7)
+    coord, atype = coord.cpu().numpy(), atype.cpu().numpy()
+    energy, force, _ = before.eval(coord, None, atype, mixed_type=True)
+    new_energy, new_force, _ = after.eval(coord, None, atype, mixed_type=True)
+    old_bias = model.get_out_bias().detach().cpu().numpy().reshape(-1)
+    # The padding atoms, of type -1, carry no bias.
+    delta = np.asarray(bias) - old_bias
+    shift = np.where(atype >= 0, delta[atype], 0.0).sum(axis=-1)
+    np.testing.assert_allclose(new_energy.reshape(-1), energy.reshape(-1) + shift)
+    np.testing.assert_allclose(new_force, force, rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize(
+    ("kernel", "source"), [("les", None), ("les", "neutral"), ("sog", "charge_spin")]
+)
+def test_coordinate_evaluation_matches_pair_graph(
+    kernel: str, source: str | None
+) -> None:
+    """The evaluation from node positions matches the pair-graph evaluation.
+
+    It returns the atomic energy, the derivative of the energy with respect to
+    the descriptor through the constrained charges, and the force at a fixed
+    descriptor, all of which the pair graph yields by automatic
+    differentiation. Two frames, an uncharged node and trailing padding nodes
+    exercise the frame sums of the constraint.
+    """
+    from deepmd.pt_expt.kernels.long_range import (
+        op_available,
+    )
+
+    if not op_available():
+        pytest.skip("the long-range operator is not available")
+    config = _config(kernel, constraint=None if source is None else {"source": source})
+    if source == "charge_spin":
+        config["descriptor"]["add_chg_spin_ebd"] = True
+    long_range = get_model(config).atomic_model.long_range.eval()
+    if source == "charge_spin":
+        long_range.fold_charge_state([2.0, 1.0])
+    rng = np.random.default_rng(7)
+    n_node = torch.tensor([9, 6], device=env.DEVICE)
+    node_mask = torch.ones(17, dtype=torch.bool, device=env.DEVICE)
+    node_mask[4] = False
+    node_mask[15:] = False
+    coord = _tensor(rng.uniform(0.0, 6.0, (17, 3)))
+    descriptor = _tensor(rng.normal(size=(17, long_range.charge_net.dim_descrpt)))
+    atype = _tensor(rng.integers(0, 2, 17), torch.int64)
+
+    atom_energy, descriptor_gradient, force, _, _ = long_range.energy_force_from_coord(
+        descriptor, atype, node_mask, n_node, coord
+    )
+
+    leaf_descriptor = descriptor.clone().requires_grad_(True)
+    leaf_coord = coord.clone().requires_grad_(True)
+    pairs = build_pair_graph(leaf_coord, n_node, atype, node_mask)
+    expected = long_range.call_graph(
+        leaf_descriptor, atype, node_mask, n_node, {ALL_PAIRS: pairs}
+    )
+    expected_descriptor_gradient, expected_coord_gradient = torch.autograd.grad(
+        expected.sum(), (leaf_descriptor, leaf_coord)
+    )
+    for actual, target in (
+        (atom_energy, expected),
+        (descriptor_gradient, expected_descriptor_gradient),
+        (force, -expected_coord_gradient),
+    ):
+        target = target.detach()
+        scale = target.abs().max().item()
+        assert scale > 0.0
+        np.testing.assert_allclose(
+            actual.detach().to(torch.float64).cpu().numpy(),
+            target.cpu().numpy(),
+            rtol=0.0,
+            atol=2e-5 * scale,
+        )
+
+
+@pytest.mark.parametrize(
+    ("kernel", "source"), [("les", "charge_spin"), ("sog", "neutral")]
+)
+def test_compressed_frozen_model_matches_uncompressed(
+    tmp_path, kernel: str, source: str
+) -> None:
+    """A compressed model keeps its long-range module through the freeze.
+
+    On a CUDA device the freeze selects the compact canonical lower, whose
+    long-range operator reads the positions of the atoms; on the CPU the
+    graph lower reads the pair lists. A constraint that reads the total
+    charge from the charge state is folded with the descriptor and follows
+    the state the deployment serves. Both match the uncompressed model.
+    """
+    config = _config(kernel, constraint={"source": source})
+    config["descriptor"].update(n_radial=8, precision="float32")
+    if source == "charge_spin":
+        config["descriptor"].update(add_chg_spin_ebd=True, default_chg_spin=[0.0, 1.0])
+    config["fitting_net"].update(
+        neuron=[32, 32], activation_function="silu", resnet_dt=False
+    )
+    config["fitting_net"]["precision"] = "float32"
+    config["long_range"]["charge_net"] = {"neuron": [16, 16], "seed": 23}
+    model = get_model(config).to(env.DEVICE).eval()
+    generator = torch.Generator().manual_seed(5)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            noise = torch.randn(parameter.shape, generator=generator)
+            parameter.add_(0.05 * noise.to(parameter))
+    reference = copy.deepcopy(model)
+    model.atomic_model.enable_compression(0.5)
+    frozen = tmp_path / f"long_range_compressed_{kernel}.pt2"
+    deserialize_to_file(
+        str(frozen),
+        {"model": model.serialize(), "model_def_script": config},
+        do_atomic_virial=True,
+        lower_kind="auto",
+    )
+    dp = DeepPot(str(frozen), auto_batch_size=False)
+    expected_lower = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph"
+    assert dp.deep_eval.metadata["lower_input_kind"] == expected_lower
+    if expected_lower == "dpa4c_canonical":
+        # The fused long-range operator serves the charge network of one
+        # channel and reads the source CSR for the per-atom virial alone; the
+        # composition of operators that serves three channels reads it always.
+        assert dp.deep_eval.metadata["canonical_requires_source_csr"] == (
+            CHANNELS[kernel] > 1
+        )
+    rng = np.random.default_rng(3)
+    coord = rng.uniform(0.0, 7.0, (1, 40, 3)) + 100.0
+    atype = rng.integers(0, 2, (1, 40))
+    names = ("energy", "force", "virial", "atom_energy", "atom_virial")
+    # The tables of the compressed descriptor approximate its networks. On the
+    # graph lower the charges of this model amplify that error at a charged
+    # state, to about 1e-4 of the scale at [1, 2]; the compact canonical lower
+    # stays within a few 1e-6.
+    tolerance = 2e-5 if expected_lower == "dpa4c_canonical" else 5e-4
+    states = [[0.0, 1.0], [1.0, 2.0], [0.0, 1.0]] if source == "charge_spin" else [None]
+    for state in states:
+        condition = {} if state is None else {"charge_spin": np.array([state])}
+        expected = reference(
+            _tensor(coord),
+            _tensor(atype, torch.int64),
+            do_atomic_virial=True,
+            **{key: _tensor(value) for key, value in condition.items()},
+        )
+        actual = dp.eval(coord, None, atype, atomic=True, **condition)
+        for name, value in zip(names, actual, strict=True):
+            target = expected[name].detach().cpu().numpy().reshape(value.shape)
+            scale = np.abs(target).max()
+            assert np.abs(value - target).max() <= tolerance * scale + 1e-5, (
+                name,
+                state,
+            )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+@REQUIRES_SUPPORTED_COMPILE
+def test_compiled_matches_eager(kernel: str) -> None:
+    """The compiled graph lower evaluates the complete model, train and eval.
+
+    In training the gradient of an energy and force loss, whose force part is
+    a second-order derivative of the energy, reaches the long-range
+    parameters as in eager execution.
+    """
+    model = _model(kernel)
+    coord, atype = _batch()
+    parameters = list(model.atomic_model.long_range.parameters())
+
+    def loss(out: dict) -> torch.Tensor:
+        return out["energy"].square().sum() + out["force"].square().sum()
+
+    for training in (True, False):
+        model.train(training)
+        compiled = _CompiledModel(
+            model, _get_model_structure_key(model), compile_eval=True
+        )
+        compiled.train(training)
+        reference = model(coord, atype)
+        result = compiled(coord, atype)
+        assert compiled._compiled_lower_by_mode, "the compiled graph was not used"
+        torch.testing.assert_close(result["energy"], reference["energy"])
+        torch.testing.assert_close(result["force"], reference["force"])
+        torch.testing.assert_close(result["virial"], reference["virial"])
+        if training:
+            expected = torch.autograd.grad(loss(reference), parameters)
+            actual = torch.autograd.grad(loss(result), parameters)
+            for expected_grad, actual_grad in zip(expected, actual, strict=True):
+                torch.testing.assert_close(actual_grad, expected_grad)
+
+
+def test_unsupported_combinations_fail() -> None:
+    """Unsupported inputs, routes, combinations and descriptors are rejected."""
+    model = _model("les")
+    coord, atype = _batch()
+    box = 10.0 * torch.eye(3, dtype=torch.float64, device=env.DEVICE).reshape(1, 9)
+    with pytest.raises(NotImplementedError, match="non-periodic"):
+        model(coord[:1], atype[:1], box=box)
+    nlist = torch.full((1, 6, 4), -1, dtype=torch.int64, device=env.DEVICE)
+    with pytest.raises(NotImplementedError, match="NeighborGraph"):
+        model.atomic_model.forward_atomic(coord[:1], atype[:1], nlist)
+    config = _config("les")
+    config["spin"] = {"use_spin": [True, False], "scheme": "native"}
+    with pytest.raises(NotImplementedError):
+        get_model(config)
+    config = _config("les")
+    config["use_srtab"] = "table.txt"
+    config["sw_rmin"], config["sw_rmax"] = 0.5, 1.0
+    with pytest.raises(NotImplementedError, match="ZBL"):
+        get_model(config)
+    config = _config("les")
+    config["type_map"] = ["O", "H"]
+    config["bridging_method"] = "ZBL"
+    with pytest.raises(NotImplementedError, match="bridging"):
+        get_model(config)
+    config = _config("les")
+    config["atom_exclude_types"] = [1]
+    with pytest.raises(NotImplementedError, match="atom_exclude_types"):
+        get_model(config)
+    child = _config("les")
+    child["type"] = "standard"
+    with pytest.raises(NotImplementedError, match="atom_exclude_types"):
+        get_model(
+            {
+                "type": "linear_ener",
+                "type_map": child["type_map"],
+                "weights": "sum",
+                "models": [child],
+                "atom_exclude_types": [1],
+            }
+        )
+    config = _config("les")
+    config["descriptor"] = {
+        "type": "se_e2_a",
+        "rcut": 3.0,
+        "rcut_smth": 2.0,
+        "sel": [8, 8],
+    }
+    with pytest.raises(NotImplementedError, match="DPA4C"):
+        get_model(config)
+
+
+def test_compressing_the_descriptor_alone_is_rejected() -> None:
+    """A folded descriptor and a constraint reading the condition are rejected."""
+    config = _config("les", constraint={"source": "charge_spin"})
+    config["descriptor"].update(
+        n_radial=8,
+        precision="float32",
+        add_chg_spin_ebd=True,
+        default_chg_spin=[0.0, 1.0],
+    )
+    config["fitting_net"]["precision"] = "float32"
+    model = get_model(config).to(env.DEVICE).eval()
+    model.atomic_model.descriptor.enable_compression(0.5)
+    coord, atype = _batch()
+    with pytest.raises(ValueError, match="folds the charge/spin condition"):
+        model(coord[:1], atype[:1])
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_dpmodel_reference_matches(kernel: str) -> None:
+    """The dpmodel reference evaluates the same energy from the portable record."""
+    from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP
+
+    model = _model(kernel)
+    coord, atype = _batch()
+    expected = model(coord, atype)
+    reference = BaseModelDP.deserialize(model.serialize())
+    ret = reference.call_common(
+        coord.cpu().numpy(), atype.cpu().numpy(), neighbor_graph_method="dense"
+    )
+    np.testing.assert_allclose(
+        ret["energy_redu"].reshape(-1),
+        expected["energy"].detach().cpu().numpy().reshape(-1),
+        rtol=1e-12,
+        atol=1e-12,
+    )
+    np.testing.assert_allclose(
+        ret["energy"].reshape(3, 6),
+        expected["atom_energy"].detach().cpu().numpy().reshape(3, 6),
+        rtol=1e-12,
+        atol=1e-12,
+    )
+
+
+def test_pair_graph_skips_padding() -> None:
+    """The pair topology joins real atoms of one frame only and ends in masked guards."""
+    coord = torch.arange(15, dtype=torch.float64).reshape(5, 3)
+    pairs = build_pair_graph(
+        coord,
+        torch.tensor([3, 2]),
+        torch.tensor([0, -1, 1, 0, 1]),
+        torch.tensor([True, False, True, True, True]),
+    )
+    real = pairs.index[:, pairs.mask]
+    assert real.tolist() == [[0, 3], [2, 4]]
+    assert (~pairs.mask).sum() == 2
+
+
+# ---------------------------------------------------------------------------
+# Periodic (k-space) evaluation
+# ---------------------------------------------------------------------------
+
+
+def _periodic_model(kernel: str, **kwargs: object) -> torch.nn.Module:
+    """A periodic model over the explicit wave numbers of every frame, ``kspace_tolerance`` in ``kwargs``."""
+    long_range = {
+        key: kwargs.pop(key) for key in ("kspace_tolerance",) if key in kwargs
+    }
+    config = _config(kernel, **kwargs)
+    config["long_range"]["periodic"] = True
+    config["long_range"]["k_grid_mode"] = "ragged"
+    config["long_range"].update(long_range)
+    return get_model(config).to(env.DEVICE).eval()
+
+
+def _periodic_case() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """A compact cluster (all atoms within rcut of each other) in a large triclinic cell."""
+    box = np.array([[9.0, 0.4, 0.0], [0.0, 8.5, 0.3], [0.0, 0.0, 9.5]])
+    rng = np.random.default_rng(7)
+    coord = 3.0 + rng.uniform(0.0, 1.4, (6, 3))
+    atype = np.array([0, 1, 0, 1, 1, 0], dtype=np.int64)
+    return coord[None], atype[None], box[None]
+
+
+def _kspace_reference(
+    coord: np.ndarray,
+    charge: np.ndarray,
+    box: np.ndarray,
+    kernel: str,
+    tolerance: float = 1e-8,
+) -> float:
+    """Independent k-space evaluation of the periodic long-range energy, in eV.
+
+    E = k_e [ 1/(2V) sum_{k != 0} K_hat(k) |rho(k)|^2 - K(0)/2 sum_i |q_i|^2 ],
+    truncated where the weight of the slowest-decaying mode falls below
+    ``tolerance``, the ``kspace_tolerance`` of the module.
+    """
+    rec = 2.0 * np.pi * np.linalg.inv(box)  # columns g_d: k(n) = rec @ n
+    if kernel == "les":
+        k_max = 2.0 * 0.8 * np.sqrt(np.log(1.0 / tolerance))
+    else:
+        k_max = np.sqrt(2.0 * np.log(1.0 / tolerance)) / 0.8
+    n_d = np.ceil(k_max * np.linalg.norm(box, axis=1) / (2 * np.pi)).astype(int)
+    grids = np.meshgrid(
+        *[np.arange(-n, n + 1, dtype=float) for n in n_d], indexing="ij"
+    )
+    n_all = np.stack(grids, axis=-1).reshape(-1, 3)
+    k = n_all @ rec.T
+    k2 = np.sum(k * k, axis=-1)
+    keep = (k2 > 0) & (k2 <= k_max * k_max)
+    k, k2 = k[keep], k2[keep]
+    phase = k @ coord.T  # (nk, N)
+    rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge  # (nk, C)
+    s_sq = np.abs(rho) ** 2
+    s_sq = s_sq.sum(axis=-1)
+    if kernel == "les":
+        k_hat = 4.0 * np.pi * np.exp(-k2 / (4.0 * 0.8**2)) / k2
+        k0 = 2.0 * 0.8 / np.sqrt(np.pi)
+    else:
+        amplitude = np.array([0.7, -0.2, 0.1])
+        width = np.array([0.8, 1.7, 3.2])
+        k_hat = np.sum(
+            amplitude[None, :]
+            * (2.0 * np.pi) ** 1.5
+            * width[None, :] ** 3
+            * np.exp(-0.5 * k2[:, None] * width[None, :] ** 2),
+            axis=-1,
+        )
+        k0 = amplitude.sum()
+    volume = abs(np.linalg.det(box))
+    return float(
+        COULOMB_CONSTANT
+        * ((k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum())
+    )
+
+
+@pytest.mark.parametrize("tolerance", [None, 1e-4])
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_energy_matches_kspace_reference(
+    kernel: str, tolerance: float | None
+) -> None:
+    """With fixed charges the periodic energy is the documented k-space sum.
+
+    The sum ends at the wave number of ``kspace_tolerance``, the default or a
+    configured one.
+    """
+    options = {} if tolerance is None else {"kspace_tolerance": tolerance}
+    model = _periodic_model(kernel, constraint=None, **options)
+    _set_constant_charges(model, [1.0, -0.5])
+    coord, atype, box = _periodic_case()
+    out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box))
+    charge = np.zeros((6, CHANNELS[kernel]))
+    charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0]
+    expected = _kspace_reference(coord[0], charge, box[0], kernel, tolerance or 1e-8)
+    if tolerance is not None:
+        # The configured truncation differs measurably from the default one.
+        default = _kspace_reference(coord[0], charge, box[0], kernel)
+        assert abs(expected - default) > 1e-6, (expected, default)
+    np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-9, atol=1e-9)
+    np.testing.assert_allclose(
+        out["atom_energy"].sum().item(), expected, rtol=1e-9, atol=1e-9
+    )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_force_virial_hessian_match_finite_differences(kernel: str) -> None:
+    """Force, strain virial and Hessian derive from the complete periodic energy."""
+    model = _periodic_model(kernel)
+    coord, atype, box = _periodic_case()
+    coord_t, atype_t, box_t = (
+        _tensor(coord),
+        _tensor(atype, torch.int64),
+        _tensor(box),
+    )
+    out = model(coord_t, atype_t, box=box_t)
+
+    def energy(c: np.ndarray, b: np.ndarray) -> float:
+        return model(_tensor(c), atype_t, box=_tensor(b))["energy"].item()
+
+    h = 1e-5
+    force_fd = np.zeros((6, 3))
+    for i in range(6):
+        for d in range(3):
+            plus, minus = coord.copy(), coord.copy()
+            plus[0, i, d] += h
+            minus[0, i, d] -= h
+            force_fd[i, d] = -(energy(plus, box) - energy(minus, box)) / (2 * h)
+    np.testing.assert_allclose(
+        out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-6
+    )
+
+    # The virial is -dE/d(strain) at zero strain, for r -> (I + eps) r and
+    # B -> B (I + eps)^T (the strain family of the k-space strain leaf).
+    virial_fd = np.zeros((3, 3))
+    for a in range(3):
+        for b in range(3):
+            strain = np.zeros((3, 3))
+            strain[a, b] = h
+            plus = coord @ (np.eye(3) + strain).T
+            minus = coord @ (np.eye(3) - strain).T
+            box_plus = box @ (np.eye(3) + strain).T
+            box_minus = box @ (np.eye(3) - strain).T
+            virial_fd[a, b] = -(energy(plus, box_plus) - energy(minus, box_minus)) / (
+                2 * h
+            )
+    np.testing.assert_allclose(
+        out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-5
+    )
+
+    hessian_model = _periodic_model(kernel)
+    hessian_model.load_state_dict(model.state_dict())
+    hessian_model.enable_hessian()
+    hessian = (
+        hessian_model(coord_t, atype_t, box=box_t)["hessian"].detach().cpu().numpy()
+    )
+    hessian = hessian.reshape(18, 18)
+    hessian_fd = np.zeros((18, 18))
+    for kk in range(18):
+        plus, minus = coord.copy(), coord.copy()
+        plus.reshape(-1)[kk] += h
+        minus.reshape(-1)[kk] -= h
+        f_plus = (
+            model(_tensor(plus), atype_t, box=box_t)["force"].detach().cpu().numpy()
+        )
+        f_minus = (
+            model(_tensor(minus), atype_t, box=box_t)["force"].detach().cpu().numpy()
+        )
+        hessian_fd[:, kk] = -(f_plus - f_minus).reshape(-1) / (2 * h)
+    np.testing.assert_allclose(hessian, hessian_fd, atol=1e-4)
+    np.testing.assert_allclose(hessian, hessian.T, atol=1e-10)
+
+
+@pytest.mark.parametrize("grid_mode", ["ragged", "pme"])
+def test_trace_samples_give_every_axis_its_own_length(grid_mode: str) -> None:
+    """No two dynamic axes of the trace-time topologies share a length, nor take 2 or 3."""
+    from deepmd.pt_expt.utils.serialization import (
+        build_synthetic_topologies,
+    )
+
+    model = (
+        _periodic_model("les") if grid_mode == "ragged" else _periodic_pme_model("les")
+    )
+    nframes, n_node, n_pair = 5, 7, 11
+    topologies = build_synthetic_topologies(
+        model, n_node, n_pair, dtype=torch.float64, device=env.DEVICE, nframes=nframes
+    )
+    lengths = {("", "frame"): nframes, ("", "node"): n_node}
+    for name, topology in topologies.items():
+        for field, axes in type(topology).AXES.items():
+            for axis, label in axes.items():
+                key = ("", label) if label in ("frame", "node") else (name, label)
+                size = getattr(topology, field).shape[axis]
+                assert lengths.setdefault(key, size) == size, (name, field, label)
+    sizes = list(lengths.values())
+    assert len(set(sizes)) == len(sizes), lengths
+    assert not {2, 3} & set(sizes), lengths
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+@REQUIRES_SUPPORTED_COMPILE
+def test_periodic_compiled_matches_eager(kernel: str) -> None:
+    """The compiled graph lower evaluates the periodic model like the eager one."""
+    model = _periodic_model(kernel)
+    coord, atype, box = _periodic_case()
+    coord_t, atype_t, box_t = (
+        _tensor(coord),
+        _tensor(atype, torch.int64),
+        _tensor(box),
+    )
+    model.train()
+    runner = _CompiledModel(model, _get_model_structure_key(model))
+    runner.train()
+    reference = model.eval()(coord_t, atype_t, box=box_t)
+    model.train()
+    out = runner(coord_t, atype_t, box=box_t)
+    for key in ("energy", "force", "virial"):
+        torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8)
+    # the backward reaches the charge network and the kernel through the
+    # compiled graph
+    loss = out["energy"].square().sum() + out["force"].square().sum()
+    loss.backward()
+    grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()]
+    assert grads and all(grad is not None for grad in grads)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_model_serves_non_periodic_batches(kernel: str) -> None:
+    """A periodic-enabled model keeps its non-periodic evaluation."""
+    model = _periodic_model(kernel)
+    plain = _model(kernel)
+    plain.load_state_dict(model.state_dict())
+    coord, atype = _batch()
+    periodic_enabled = model(coord, atype)
+    plain_out = plain(coord, atype)
+    torch.testing.assert_close(periodic_enabled["energy"], plain_out["energy"])
+    torch.testing.assert_close(periodic_enabled["force"], plain_out["force"])
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_frames_are_independent(kernel: str) -> None:
+    """Each frame's k grid and energy are independent of the rest of the batch."""
+    model = _periodic_model(kernel)
+    coord, atype, box = _periodic_case()
+    coord2 = np.concatenate([coord, coord + 0.1], axis=0)
+    atype2 = np.concatenate([atype, atype], axis=0)
+    box2 = np.concatenate([box, box * 1.1], axis=0)
+    batched = model(_tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2))
+    for frame in range(2):
+        alone = model(
+            _tensor(coord2[frame : frame + 1]),
+            _tensor(atype2[frame : frame + 1], torch.int64),
+            box=_tensor(box2[frame : frame + 1]),
+        )
+        torch.testing.assert_close(
+            alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10
+        )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_serialization_round_trip(kernel: str) -> None:
+    """The record keeps the periodic flag and evaluates identically."""
+    from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP
+
+    model = _periodic_model(kernel)
+    assert model.atomic_model.required_topologies() == ("all_pairs", "k_grid")
+    record = model.serialize()
+    assert record["long_range"]["periodic"] is True
+    reference = BaseModelDP.deserialize(record)
+    assert reference.atomic_model.required_topologies() == ("all_pairs", "k_grid")
+    coord, atype, box = _periodic_case()
+    expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box))
+    ret = reference.call_common(
+        coord,
+        atype,
+        box,
+        neighbor_graph_method="dense",
+    )
+    np.testing.assert_allclose(
+        ret["energy_redu"].reshape(-1),
+        expected["energy"].detach().cpu().numpy().reshape(-1),
+        rtol=1e-9,
+        atol=1e-9,
+    )
+
+
+def _periodic_frames() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Two frames of the periodic case, the second strained and shifted."""
+    coord, atype, box = _periodic_case()
+    strain = np.array([[1.04, 0.02, 0.0], [0.0, 0.97, 0.01], [0.03, 0.0, 1.02]])
+    coord2 = np.concatenate([coord, coord @ strain.T + 0.3], axis=0)
+    box2 = np.concatenate([box, box @ strain.T], axis=0)
+    return coord2, np.concatenate([atype, atype], axis=0), box2
+
+
+@pytest.mark.parametrize(
+    ("kernel", "grid_mode", "suffix"),
+    [("les", "ragged", ".pt2"), ("sog", "pme", ".pte"), ("sog", "fixed", ".pt2")],
+)
+def test_periodic_frozen_model_matches_mesh_evaluation(
+    tmp_path, kernel: str, grid_mode: str, suffix: str
+) -> None:
+    """A frozen periodic model evaluates every periodic frame on the mesh of its cell.
+
+    Whatever its training mode, the frozen model reads the mesh of the cells
+    of a batch as inputs of dynamic size, and matches the model evaluated on
+    the mesh; a batch without a cell keeps the direct sum.
+    """
+    if suffix == ".pt2" and not mesh_freeze_supported():
+        pytest.skip(f"torch {torch.__version__} cannot package the mesh evaluation")
+    config = _config(kernel)
+    config["long_range"]["periodic"] = True
+    config["long_range"]["k_grid_mode"] = grid_mode
+    if grid_mode == "fixed":
+        # Bounds that hold every wave number of the cells within the cutoff.
+        config["long_range"]["k_grid_bounds"] = [14, 14, 14]
+    model = get_model(config).to(env.DEVICE).eval()
+    frozen = tmp_path / f"periodic_{kernel}_{grid_mode}{suffix}"
+    deserialize_to_file(
+        str(frozen),
+        {"model": model.serialize(), "model_def_script": config},
+        do_atomic_virial=True,
+        lower_kind="graph",
+    )
+    dp = DeepPot(str(frozen), auto_batch_size=False)
+    assert dp.deep_eval.metadata["topologies"] == ["all_pairs", "pme_mesh"]
+    mesh_model = copy.deepcopy(model)
+    mesh_model.atomic_model.long_range.use_mesh_evaluation()
+    names = ("energy", "force", "virial", "atom_energy", "atom_virial")
+    coord, atype, box = _periodic_frames()
+    for frame in (slice(0, 1), slice(0, 2)):
+        expected = mesh_model(
+            _tensor(coord[frame]),
+            _tensor(atype[frame], torch.int64),
+            box=_tensor(box[frame]),
+            do_atomic_virial=True,
+        )
+        actual = dp.eval(
+            coord[frame], box[frame], atype[frame], atomic=True, mixed_type=True
+        )
+        for name, value in zip(names, actual, strict=True):
+            np.testing.assert_allclose(
+                value,
+                expected[name].detach().cpu().numpy().reshape(value.shape),
+                rtol=1e-10,
+                atol=1e-10,
+                err_msg=name,
+            )
+    # The mesh approximates the sum of the training mode.
+    trained = model(
+        _tensor(coord[:1]), _tensor(atype[:1], torch.int64), box=_tensor(box[:1])
+    )
+    np.testing.assert_allclose(
+        dp.eval(coord[:1], box[:1], atype[:1], mixed_type=True)[0],
+        trained["energy"].detach().cpu().numpy(),
+        rtol=1e-5,
+    )
+    cluster, cluster_type = _batch()
+    expected = model(cluster, cluster_type)
+    actual = dp.eval(
+        cluster.cpu().numpy(), None, cluster_type.cpu().numpy(), mixed_type=True
+    )
+    np.testing.assert_allclose(
+        actual[0], expected["energy"].detach().cpu().numpy(), rtol=1e-10, atol=1e-10
+    )
+
+
+@pytest.mark.parametrize(("kernel", "source"), [("les", "neutral"), ("sog", None)])
+def test_periodic_compressed_frozen_model_matches_uncompressed(
+    tmp_path, kernel: str, source: str | None
+) -> None:
+    """A compressed periodic model keeps the mesh evaluation through the freeze.
+
+    On a CUDA device the compact canonical lower evaluates every periodic
+    frame on the mesh of its cell with the long-range operator; on the CPU
+    the graph lower reads the mesh. Both match the uncompressed model on the
+    mesh, for a periodic frame and for a cluster.
+    """
+    if env.DEVICE.type != "cuda" and not mesh_freeze_supported():
+        pytest.skip(f"torch {torch.__version__} cannot package the mesh evaluation")
+    config = _config(kernel, constraint=None if source is None else {"source": source})
+    config["descriptor"].update(n_radial=8, precision="float32")
+    config["fitting_net"].update(
+        neuron=[32, 32], activation_function="silu", resnet_dt=False
+    )
+    config["fitting_net"]["precision"] = "float32"
+    config["long_range"]["charge_net"] = {"neuron": [16, 16], "seed": 23}
+    config["long_range"]["periodic"] = True
+    model = get_model(config).to(env.DEVICE).eval()
+    generator = torch.Generator().manual_seed(5)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            noise = torch.randn(parameter.shape, generator=generator)
+            parameter.add_(0.05 * noise.to(parameter))
+    reference = copy.deepcopy(model)
+    reference.atomic_model.long_range.use_mesh_evaluation()
+    model.atomic_model.enable_compression(0.5)
+    frozen = tmp_path / f"periodic_compressed_{kernel}.pt2"
+    deserialize_to_file(
+        str(frozen),
+        {"model": model.serialize(), "model_def_script": config},
+        do_atomic_virial=True,
+        lower_kind="auto",
+    )
+    dp = DeepPot(str(frozen), auto_batch_size=False)
+    expected_lower = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph"
+    assert dp.deep_eval.metadata["lower_input_kind"] == expected_lower
+    if expected_lower == "dpa4c_canonical":
+        assert dp.deep_eval.metadata["canonical_requires_source_csr"] == (
+            CHANNELS[kernel] > 1
+        )
+    names = ("energy", "force", "virial", "atom_energy", "atom_virial")
+    rng = np.random.default_rng(3)
+    box = np.array([[[12.0, 0.0, 0.0], [0.8, 11.0, 0.0], [-0.5, 0.6, 13.0]]])
+    periodic = rng.uniform(0.0, 1.0, (1, 40, 3)) @ box[0]
+    cluster = rng.uniform(0.0, 7.0, (1, 40, 3)) + 100.0
+    atype = rng.integers(0, 2, (1, 40))
+    for coord, cell in ((periodic, box), (cluster, None)):
+        expected = reference(
+            _tensor(coord),
+            _tensor(atype, torch.int64),
+            box=None if cell is None else _tensor(cell),
+            do_atomic_virial=True,
+        )
+        actual = dp.eval(coord, cell, atype, atomic=True)
+        for name, value in zip(names, actual, strict=True):
+            target = expected[name].detach().cpu().numpy().reshape(value.shape)
+            scale = np.abs(target).max()
+            # The tables of the compressed descriptor approximate its
+            # networks, to a few 1e-6 of the scale of every output.
+            assert np.abs(value - target).max() <= 2e-5 * scale + 1e-5, (
+                name,
+                cell is None,
+                np.abs(value - target).max(),
+                scale,
+            )
+    # A cell without volume marks no system, on either lower: the graph lower
+    # meets it in the neighbor search, the compact lower at its input.
+    with pytest.raises((ValueError, RuntimeError), match=r"no volume|singular"):
+        dp.eval(periodic, np.zeros((1, 9)), atype)
+
+
+def _canonical_inputs(
+    model: torch.nn.Module,
+    coord: np.ndarray,
+    atype: np.ndarray,
+    box: np.ndarray | None,
+) -> tuple:
+    """The compact canonical graph of the frames, as one rank builds it, and the node inputs of the module."""
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        NeighborGraph,
+        build_neighbor_graph,
+    )
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_graph_from_neighbor_graph,
+    )
+
+    nf = atype.shape[0]
+    cells = None if box is None else box.reshape(nf, 9)
+    graph = build_neighbor_graph(
+        coord, atype, cells, model.get_rcut(), canonicalize=True
+    )
+    n_node = torch.as_tensor(graph.n_node, dtype=torch.int64, device=env.DEVICE)
+    graph = NeighborGraph(
+        n_node=n_node,
+        n_local=n_node,
+        edge_index=torch.as_tensor(graph.edge_index, device=env.DEVICE),
+        edge_vec=torch.as_tensor(
+            graph.edge_vec, dtype=torch.float32, device=env.DEVICE
+        ),
+        edge_mask=torch.as_tensor(graph.edge_mask, device=env.DEVICE),
+        destination_order=torch.as_tensor(graph.destination_order, device=env.DEVICE),
+        destination_row_ptr=torch.as_tensor(
+            graph.destination_row_ptr, dtype=torch.int64, device=env.DEVICE
+        ),
+        source_order=torch.as_tensor(graph.source_order, device=env.DEVICE),
+        source_row_ptr=torch.as_tensor(
+            graph.source_row_ptr, dtype=torch.int64, device=env.DEVICE
+        ),
+        destination_sorted=bool(graph.destination_sorted),
+    )
+    compact = canonical_graph_from_neighbor_graph(graph)
+    return (
+        _tensor(atype.reshape(-1), torch.int64),
+        compact,
+        _tensor(coord.reshape(-1, 3)),
+        None if cells is None else torch.as_tensor(cells, dtype=torch.float64),
+    )
+
+
+# The layouts of the fused parity test: the width of the energy fitting and the
+# hidden layers of the charge network, which select the per-layer kernels at
+# widths of 32 and 16 and the fused network at widths of 96 and 64, and the
+# frames of the batch.
+FUSED_LAYOUTS = {
+    "per_layer": {"width": 32, "charge": [16, 16, 16], "frames": 2},
+    "fused": {"width": 96, "charge": [64, 64], "frames": 1},
+}
+
+
+@pytest.mark.skipif(
+    env.DEVICE.type != "cuda", reason="the compact canonical lower is CUDA only"
+)
+@pytest.mark.parametrize("tile", [0, 16])
+@pytest.mark.parametrize("layout", ["per_layer", "fused"])
+@pytest.mark.parametrize("periodic", [False, True])
+@pytest.mark.parametrize(("kernel", "source"), [("les", "neutral"), ("sog", None)])
+def test_fused_operator_matches_composition(
+    monkeypatch,
+    kernel: str,
+    source: str | None,
+    periodic: bool,
+    layout: str,
+    tile: int,
+) -> None:
+    """The fused long-range operator evaluates a compressed model like the composition.
+
+    With a charge network of one channel, the compact canonical lower
+    evaluates the descriptor, both networks, the long range and the force
+    assembly in one operator; the composition of operators, which serves the
+    other charge networks, is the reference. The energy fitting and the
+    charge network take the per-layer kernels or the fused network, over the
+    whole node axis or in runs of 16 nodes with a shorter last one, with and
+    without the per-atom virial, on two frames, the second with fewer real
+    atoms, or on one.
+    """
+    from deepmd.pt_expt.kernels.dpa4c import (
+        canonical_long_range,
+    )
+    from deepmd.pt_expt.kernels.graph_fitting import (
+        fitting_operator_arguments,
+    )
+
+    shape = FUSED_LAYOUTS[layout]
+    monkeypatch.setenv("DP_NODE_TILE", str(tile))
+    config = _config(kernel, constraint=None if source is None else {"source": source})
+    config["descriptor"].update(n_radial=8, precision="float32")
+    config["fitting_net"].update(
+        neuron=[shape["width"]] * 2, activation_function="silu", resnet_dt=False
+    )
+    config["fitting_net"]["precision"] = "float32"
+    config["long_range"].update(
+        channels=1,
+        charge_net={"neuron": shape["charge"], "seed": 23},
+        periodic=periodic,
+    )
+    if kernel == "sog":
+        config["long_range"]["kernel"] = {"type": "sog", "n_gaussian": 4}
+    model = get_model(config).to(env.DEVICE).eval()
+    generator = torch.Generator().manual_seed(5)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            noise = torch.randn(parameter.shape, generator=generator)
+            parameter.add_(0.05 * noise.to(parameter))
+    if periodic:
+        model.atomic_model.long_range.use_mesh_evaluation()
+    model.atomic_model.enable_compression(0.5)
+    assert canonical_long_range.fused_long_range_eligible(model.atomic_model.long_range)
+    frames = shape["frames"]
+    rng = np.random.default_rng(3)
+    box = np.array([[[12.0, 0.0, 0.0], [0.8, 11.0, 0.0], [-0.5, 0.6, 13.0]]] * frames)
+    coord = (
+        rng.uniform(0.0, 1.0, (frames, 30, 3)) @ box[0]
+        if periodic
+        else (rng.uniform(0.0, 7.0, (frames, 30, 3)) + 100.0)
+    )
+    atype = rng.integers(0, 2, (frames, 30))
+    atype[-1, 24:] = -1
+    atype_t, graph, coord_t, box_t = _canonical_inputs(
+        model, coord, atype, box if periodic else None
+    )
+    # The rules of the fitting routes take the network of the layout on every
+    # run of the energy fitting and of the charge network; a device on which
+    # they send the runs of the fused layout to the per-layer kernels has no
+    # fused network to compare.
+    nodes = atype_t.shape[0]
+    runs = {nodes} if tile == 0 else {tile, nodes % tile or tile}
+    networks = (
+        model.atomic_model.fitting_net,
+        model.atomic_model.long_range.charge_net,
+    )
+    weights = [fitting_operator_arguments(network).weights for network in networks]
+    routes = [
+        int(torch.ops.deepmd.graph_fitting_route_rule(ws, n))
+        for ws in weights
+        for n in runs
+    ]
+    if layout == "fused" and -1 in routes:
+        pytest.skip("the fitting rules take the per-layer kernels on this device")
+    assert all((route == -1) == (layout == "per_layer") for route in routes), routes
+
+    def evaluate(do_atomic_virial: bool) -> dict[str, torch.Tensor]:
+        return model.forward_lower_canonical_graph(
+            atype_t,
+            graph.n_node,
+            graph.n_local,
+            graph.source,
+            graph.edge_vec,
+            graph.destination_row_ptr,
+            graph.source_row_ptr,
+            graph.source_order,
+            do_atomic_virial=do_atomic_virial,
+            coord=coord_t,
+            box=box_t,
+        )
+
+    for do_atomic_virial in (True, False):
+        fused = evaluate(do_atomic_virial)
+        with monkeypatch.context() as patch:
+            patch.setattr(
+                canonical_long_range, "fused_long_range_eligible", lambda _: False
+            )
+            composed = evaluate(do_atomic_virial)
+        assert set(fused) == set(composed)
+        for key, value in composed.items():
+            # Both evaluate in float32 and sum in other orders, so they agree
+            # to a few roundings of the largest entry of each output; the
+            # long range dominates these outputs.
+            expected = value.detach().to(torch.float64)
+            error = (fused[key].detach().to(torch.float64) - expected).abs().max()
+            scale = expected.abs().max()
+            assert error <= 1e-5 * scale + 1e-6, (key, error.item(), scale.item())
+
+
+#: Bit widths of the template parameters of a fused fitting shape in its route
+#: key (FusedShape::key of graph_fitting.cu), from the lowest bit up.
+_SHAPE_KEY_FIELDS = (9, 6, 6, 6, 6, 6, 6, 6, 7, 4)
+
+
+def _fused_shape_name(key: int) -> str:
+    """Return the template arguments of the fused shape of a route key, as a trace names them."""
+    values = []
+    for bits in _SHAPE_KEY_FIELDS:
+        values.append(key & ((1 << bits) - 1))
+        key >>= bits
+    return "FusedShape<" + ", ".join(str(v) for v in values) + ">"
+
+
+def _network_routes(regime: str) -> dict:
+    """Evaluate the fused long-range operator under every route of either network.
+
+    Builds a compressed model with an energy fitting of width 96 and a charge
+    network of width 64 on a cluster of a run of fewer 32-node panels than
+    multiprocessors (``latency``) or of more (``throughput``), and evaluates
+    the operator under empty tables, the rules, under a table naming each
+    route of one network at a time, under the table that the route tuner holds
+    for the charge network, as the energy model looks it up, here a route
+    besides its rule, and, split into a full run and a shorter tail, under
+    tables that send one of them to the per-layer kernels of the charge
+    network and the other to a fused shape, and both to fused shapes. The
+    fitting kernels of an evaluation come from a kernel trace.
+
+    Returns
+    -------
+    dict
+        ``rules``, the fitting kernels of the rules; ``routes``, for every
+        untiled table, whether all six outputs equal those of the rules bit
+        for bit and its fitting kernels; ``mixed``, the same for every mixed
+        table against the rules tiled alike, with its launches of the
+        per-layer forward kernels; and ``fused_layers``, those launches with
+        both runs of the charge network on fused shapes, or ``None`` where
+        either run has no fused shape.
+    """
+    from torch.profiler import (
+        ProfilerActivity,
+        profile,
+    )
+
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+    from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import (
+        dpa4c_canonical_compress_long_range_energy_force,
+    )
+    from deepmd.pt_expt.kernels.fitting_routes import (
+        PER_LAYER,
+        ROUTE_TABLE_MAGIC,
+    )
+    from deepmd.pt_expt.kernels.graph_fitting import (
+        fitting_operator_arguments,
+    )
+    from deepmd.pt_expt.model.make_model import (
+        fused_atom_bias,
+    )
+
+    config = _config("les", constraint={"source": "neutral"})
+    config["descriptor"].update(n_radial=8, precision="float32")
+    config["fitting_net"].update(
+        neuron=[96] * 2, activation_function="silu", resnet_dt=False
+    )
+    config["fitting_net"]["precision"] = "float32"
+    config["long_range"].update(channels=1, charge_net={"neuron": [64, 64], "seed": 23})
+    model = get_model(config).to(env.DEVICE).eval()
+    generator = torch.Generator().manual_seed(5)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            noise = torch.randn(parameter.shape, generator=generator)
+            parameter.add_(0.05 * noise.to(parameter))
+    model.atomic_model.enable_compression(0.5)
+    device = env.DEVICE.index or 0
+    multiprocessors = torch.cuda.get_device_properties(env.DEVICE).multi_processor_count
+    nodes = 30 if regime == "latency" else int(32 * multiprocessors * 1.5)
+    rng = np.random.default_rng(3)
+    # A cluster of about one atom per 20 cubic angstroms.
+    side = (20.0 * nodes) ** (1.0 / 3.0)
+    coord = rng.uniform(0.0, side, (1, nodes, 3)) + 100.0
+    atype = rng.integers(0, 2, (1, nodes))
+    atype_t, graph, coord_t, _ = _canonical_inputs(model, coord, atype, None)
+    learned, _ = model.atomic_model.fused_decomposition()
+    atom_bias = fused_atom_bias(model.atomic_model, learned)
+    fingerprint = [
+        int(v) for v in torch.ops.deepmd.graph_fitting_route_fingerprint(device)
+    ]
+
+    def evaluate(routes: list[int], charge_routes: list[int]) -> tuple:
+        return dpa4c_canonical_compress_long_range_energy_force(
+            learned.descriptor,
+            learned.fitting_net,
+            learned.long_range,
+            graph,
+            atype_t,
+            coord_t,
+            None,
+            None,
+            atom_bias,
+            fitting_routes=routes,
+            charge_fitting_routes=charge_routes,
+        )
+
+    def choices(width: int, size: int = nodes) -> list[int]:
+        fused = torch.ops.deepmd.graph_fitting_route_choices(device, width, size)
+        return [PER_LAYER, *(int(key) for key in fused)]
+
+    def kernel_names(call: Any) -> list[str]:
+        with profile(activities=[ProfilerActivity.CUDA]) as trace:
+            call()
+            torch.cuda.synchronize()
+        names = [
+            event.name
+            for event in trace.events()
+            if event.device_type == torch.autograd.DeviceType.CUDA
+        ]
+        assert names, "the profiler recorded no CUDA kernel"
+        return names
+
+    def kernels(call: Any) -> list[str]:
+        names = kernel_names(call)
+        return sorted(
+            {name for name in names if "fused_network_kernel" in name}
+            | {"per-layer" for name in names if "layer_forward_kernel" in name}
+        )
+
+    def layer_launches(call: Any) -> int:
+        return sum("layer_forward_kernel" in name for name in kernel_names(call))
+
+    results = {}
+    with torch.no_grad():
+        rules = evaluate([], [])
+        tables = [
+            ("energy", c, [ROUTE_TABLE_MAGIC, *fingerprint, 1, c], [])
+            for c in choices(96)
+        ]
+        tables += [
+            ("charge", c, [], [ROUTE_TABLE_MAGIC, *fingerprint, 1, c])
+            for c in choices(64)
+        ]
+        charge_net = learned.long_range.charge_net
+        rule = int(
+            torch.ops.deepmd.graph_fitting_route_rule(
+                fitting_operator_arguments(charge_net).weights, nodes
+            )
+        )
+        measured = next(c for c in choices(64) if c != rule)
+        key = (
+            fitting_routes._fingerprint(env.DEVICE),
+            fitting_routes.network_key(charge_net),
+        )
+        table = (ROUTE_TABLE_MAGIC, *fingerprint, 1, measured)
+        fitting_routes._ROUTE_TABLES[key] = table
+        tables.append(
+            ("charge-measured", measured, [], fitting_routes.fitting_routes(charge_net))
+        )
+        for network, choice, routes, charge_routes in tables:
+            routed = evaluate(routes, charge_routes)
+            equal = all(
+                torch.equal(expected, actual)
+                for expected, actual in zip(rules, routed, strict=True)
+            )
+            launched = kernels(lambda: evaluate(routes, charge_routes))
+            results[f"{network} {choice}"] = {"equal": equal, "kernels": launched}
+        rules_kernels = kernels(lambda: evaluate([], []))
+        # One full run and a shorter tail, of which one takes the per-layer
+        # kernels of the charge network and the other a fused shape, against
+        # both on fused shapes.
+        tile = nodes // 2 + 1
+        tail = nodes - tile
+        os.environ["DP_NODE_TILE"] = str(tile)
+        tiled_rules = evaluate([], [])
+        fused_tail, fused_full = choices(64, tail)[1:2], choices(64, tile)[1:2]
+        mixed = {}
+        fused_layers = None
+
+        def two_runs(tail_route: int, full_route: int) -> list[int]:
+            return [ROUTE_TABLE_MAGIC, *fingerprint, 2, tail, tail_route, full_route]
+
+        if fused_tail and fused_full:
+            # The per-layer kernels of the energy fitting alone.
+            table = two_runs(*fused_tail, *fused_full)
+            fused_layers = layer_launches(lambda: evaluate([], table))
+            pairs = [(PER_LAYER, *fused_full), (*fused_tail, PER_LAYER)]
+            for tail_route, full_route in pairs:
+                table = two_runs(tail_route, full_route)
+                routed = evaluate([], table)
+                equal = all(
+                    torch.equal(expected, actual)
+                    for expected, actual in zip(tiled_rules, routed, strict=True)
+                )
+                mixed[f"{tail_route} {full_route}"] = {
+                    "equal": equal,
+                    "kernels": kernels(lambda: evaluate([], table)),
+                    "layers": layer_launches(lambda: evaluate([], table)),
+                }
+        del os.environ["DP_NODE_TILE"]
+    return {
+        "rules": rules_kernels,
+        "routes": results,
+        "mixed": mixed,
+        "fused_layers": fused_layers,
+    }
+
+
+@pytest.mark.skipif(
+    env.DEVICE.type != "cuda", reason="the compact canonical lower is CUDA only"
+)
+@pytest.mark.parametrize("regime", ["latency", "throughput"])  # run size regime
+def test_every_network_route_returns_the_bits_of_the_rules(regime: str) -> None:
+    """Every route of either network gives the outputs of the rules bit for bit.
+
+    The fused long-range operator evaluates the energy fitting and the charge
+    network of a module by a fitting-route table each. A table that names the
+    per-layer kernels or a fused shape for one network routes that network
+    there and leaves the other at its rules, which the fitting kernels of a
+    trace show, and every route evaluates the same contractions in the same
+    order, so all six outputs equal those of the rules, which empty tables
+    take. The table that the route tuner holds for the charge network, as the
+    energy model looks it up, routes it as well, and a tiled call whose full
+    run and tail take the per-layer kernels and a fused shape of the charge
+    network, in either order, keeps the bits of the rules. The trace runs in a
+    process of its own, untiled but for that call, since CUPTI can stop
+    delivering kernel records late in a long process. A run of fewer 32-node
+    panels than multiprocessors and one of more cover the latency and the
+    throughput shapes.
+    """
+    from deepmd.pt_expt.kernels.fitting_routes import (
+        PER_LAYER,
+    )
+
+    root = pathlib.Path(__file__).parents[__name__.count(".")]
+    code = (
+        "import importlib, json, sys\n"
+        f"sys.path[:0] = {[str(root), *sys.path]!r}\n"
+        f"module = importlib.import_module({__name__!r})\n"
+        f"print('result:', json.dumps(module._network_routes({regime!r})))\n"
+    )
+    process = subprocess.run(
+        [sys.executable, "-c", code],
+        capture_output=True,
+        text=True,
+        check=False,
+        env={**os.environ, "DP_NODE_TILE": "0"},
+    )
+    assert process.returncode == 0, process.stdout + "\n" + process.stderr
+    (report,) = [
+        line for line in process.stdout.splitlines() if line.startswith("result:")
+    ]
+    report = json.loads(report.removeprefix("result:"))
+    rules, results = report["rules"], report["routes"]
+    fused = [label for label in results if int(label.split()[1]) != PER_LAYER]
+    assert any(label.startswith("charge") for label in fused), (
+        "no fused shape of the charge network runs on this device"
+    )
+
+    def shapes(kernels: list[str], width: int) -> list[str]:
+        return [k for k in kernels if f"FusedShape<{width}," in k]
+
+    for label, result in results.items():
+        network, choice = label.split()[0], int(label.split()[1])
+        assert result["equal"], label
+        width, other = (96, 64) if network == "energy" else (64, 96)
+        own = shapes(result["kernels"], width)
+        if choice == PER_LAYER:
+            assert not own and "per-layer" in result["kernels"], (label, result)
+        else:
+            shape = _fused_shape_name(choice)
+            assert own and all(shape in k for k in own), (label, result)
+        assert shapes(result["kernels"], other) == shapes(rules, other), (
+            label,
+            result,
+            rules,
+        )
+    assert report["mixed"], "no fused shape of the charge network serves both runs"
+    for label, result in report["mixed"].items():
+        assert result["equal"], label
+        assert shapes(result["kernels"], 64), (label, result)
+        # The run of the per-layer kernels of the charge network adds layer
+        # launches to those of the energy fitting.
+        assert result["layers"] > report["fused_layers"], (label, result, report)
+
+
+@pytest.mark.skipif(
+    env.DEVICE.type != "cuda", reason="the compact canonical lower is CUDA only"
+)
+@pytest.mark.parametrize("charge_width", [64, 96])  # charge network width
+def test_the_route_tuner_measures_the_charge_network(
+    monkeypatch, charge_width: int
+) -> None:
+    """The route tuner measures the charge network of a fused module apart.
+
+    The tuner records a table for the energy fitting and one for the charge
+    network, each under its own network key, which the energy model passes to
+    the fused long-range operator; a charge network of the shape of the
+    energy fitting shares its key and its table, measured once. The
+    measurement itself is replaced by a stand-in that records which network
+    it times.
+    """
+    from deepmd.pt_expt.kernels import (
+        fitting_routes,
+    )
+    from deepmd.pt_expt.kernels.graph_fitting import (
+        fitting_operator_arguments,
+    )
+
+    config = _config("les", constraint={"source": "neutral"})
+    config["descriptor"].update(n_radial=8, precision="float32")
+    config["fitting_net"].update(
+        neuron=[96] * 2, activation_function="silu", resnet_dt=False
+    )
+    config["fitting_net"]["precision"] = "float32"
+    config["long_range"].update(
+        channels=1,
+        charge_net={"neuron": [charge_width] * 2, "seed": 23},
+    )
+    config["long_range"]["charge_net"].update(
+        activation_function="silu", resnet_dt=False
+    )
+    model = get_model(config).to(env.DEVICE).eval()
+    model.atomic_model.enable_compression(0.5)
+    learned, _ = model.atomic_model.fused_decomposition()
+    energy_key = fitting_routes.network_key(learned.fitting_net)
+    charge_key = fitting_routes.network_key(learned.long_range.charge_net)
+    assert energy_key is not None and charge_key is not None
+    assert (energy_key == charge_key) == (charge_width == 96)
+
+    measured = []
+
+    def measure(bench, ws, largest):
+        measured.append(int(ws[0].shape[1]))
+        return [], [fitting_routes.RULES]
+
+    tables: dict = {}
+    monkeypatch.setattr(fitting_routes, "_ROUTE_TABLES", tables)
+    monkeypatch.setattr(fitting_routes, "_measure", measure)
+    device = torch.device("cuda", env.DEVICE.index or 0)
+    fitting_routes.tune_fitting_routes(model, 0, device)
+    fingerprint = fitting_routes._fingerprint(device)
+    assert set(tables) == {(fingerprint, energy_key), (fingerprint, charge_key)}
+    assert sorted(measured) == sorted({96, charge_width})
+    for network in (learned.fitting_net, learned.long_range.charge_net):
+        widths = [int(w.shape[1]) for w in fitting_operator_arguments(network).weights]
+        assert widths[0] in measured
+
+
+def test_library_without_the_fused_operator_keeps_the_composition(monkeypatch) -> None:
+    """A library without the fused long-range operator evaluates the composition.
+
+    The fused operator is optional beside the short-range suite: without it,
+    a module that it would serve takes the composition of operators, and the
+    archive records that its compact lower reads the source CSR.
+    """
+    from deepmd.pt_expt.kernels.dpa4c import (
+        canonical,
+        canonical_long_range,
+    )
+    from deepmd.pt_expt.utils.serialization import (
+        _canonical_requires_source_csr,
+    )
+
+    config = _config("les")
+    config["descriptor"]["precision"] = "float32"
+    config["fitting_net"].update(
+        neuron=[32, 32], activation_function="silu", resnet_dt=False
+    )
+    config["fitting_net"]["precision"] = "float32"
+    model = get_model(config)
+    monkeypatch.setattr(
+        canonical,
+        "operator_available",
+        lambda name: name != "dpa4c_canonical_compress_long_range_energy_force",
+    )
+    assert canonical.op_available()
+    assert not canonical.long_range_op_available()
+    assert not canonical_long_range.fused_long_range_eligible(
+        model.atomic_model.long_range
+    )
+    assert _canonical_requires_source_csr(model, "dpa4c_canonical")
+
+
+# ---------------------------------------------------------------------------
+# Periodic fixed k-grid mode
+# ---------------------------------------------------------------------------
+
+
+def _periodic_fixed_model(
+    kernel: str, boxes: np.ndarray | None = None, **kwargs: object
+) -> torch.nn.Module:
+    """A periodic model in the fixed k-grid mode, sized by the given cells."""
+    config = _config(kernel, **kwargs)
+    config["long_range"]["periodic"] = True
+    config["long_range"]["k_grid_mode"] = "fixed"
+    model = get_model(config).to(env.DEVICE).eval()
+    if boxes is not None:
+        boxes = torch.as_tensor(np.asarray(boxes, dtype=np.float64).reshape(-1, 9))
+        model.atomic_model.long_range.compute_input_stats([{"box": boxes}])
+    return model
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_fixed_grid_matches_ragged(kernel: str) -> None:
+    """The fixed grid with its mask evaluates like the ragged enumeration."""
+    coord, atype, box = _periodic_case()
+    ragged = _periodic_model(kernel)
+    fixed = _periodic_fixed_model(kernel, box)
+    # the state dicts of the two grid modes differ by the fixed grid buffer
+    fixed.load_state_dict(ragged.state_dict(), strict=False)
+    coord_t, atype_t, box_t = (
+        _tensor(coord),
+        _tensor(atype, torch.int64),
+        _tensor(box),
+    )
+    expected = ragged(coord_t, atype_t, box=box_t)
+    got = fixed(coord_t, atype_t, box=box_t)
+    for key in ("energy", "force", "virial"):
+        torch.testing.assert_close(got[key], expected[key], rtol=1e-9, atol=1e-12)
+    # the gradient reaches the charge network and the kernel without NaN,
+    # though masked points carry a singular kernel transform nowhere
+    fixed.train()
+    out = fixed(coord_t, atype_t, box=box_t)
+    loss = out["energy"].square().sum() + out["force"].square().sum()
+    loss.backward()
+    grads = [
+        parameter.grad
+        for parameter in fixed.atomic_model.long_range.parameters()
+        if parameter.grad is not None
+    ]
+    assert grads and all(torch.isfinite(grad).all() for grad in grads)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_fixed_grid_mixed_cells(kernel: str) -> None:
+    """One shared grid serves a batch of different cells, each masked to its sphere."""
+    coord, atype, box = _periodic_case()
+    coord2 = np.concatenate([coord, coord + 0.1], axis=0)
+    atype2 = np.concatenate([atype, atype], axis=0)
+    box2 = np.concatenate([box, box * 1.4], axis=0)
+    fixed = _periodic_fixed_model(kernel, box2)
+    ragged = _periodic_model(kernel)
+    ragged.load_state_dict(fixed.state_dict(), strict=False)
+    coord_t, atype_t, box_t = (
+        _tensor(coord2),
+        _tensor(atype2, torch.int64),
+        _tensor(box2),
+    )
+    batched = fixed(coord_t, atype_t, box=box_t)
+    expected = ragged(coord_t, atype_t, box=box_t)
+    torch.testing.assert_close(
+        batched["energy"], expected["energy"], rtol=1e-9, atol=1e-12
+    )
+    torch.testing.assert_close(
+        batched["force"], expected["force"], rtol=1e-9, atol=1e-12
+    )
+    # the batched energy of a frame equals its standalone evaluation
+    for frame in range(2):
+        alone = fixed(
+            _tensor(coord2[frame : frame + 1]),
+            _tensor(atype2[frame : frame + 1], torch.int64),
+            box=_tensor(box2[frame : frame + 1]),
+        )
+        torch.testing.assert_close(
+            alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10
+        )
+
+
+def _kspace_reference_bounded(
+    coord: np.ndarray,
+    charge: np.ndarray,
+    box: np.ndarray,
+    kernel: str,
+    bounds: np.ndarray,
+) -> float:
+    """The k-space reference of :func:`_kspace_reference` restricted to the
+    integer indices within ``bounds`` — the truncated sum of a fixed grid.
+    """
+    rec = 2.0 * np.pi * np.linalg.inv(box)
+    if kernel == "les":
+        k_max = 2.0 * 0.8 * np.sqrt(np.log(1e8))
+    else:
+        k_max = np.sqrt(2.0 * np.log(1e8)) / 0.8
+    grids = np.meshgrid(
+        *[np.arange(-n, n + 1, dtype=float) for n in bounds], indexing="ij"
+    )
+    n_all = np.stack(grids, axis=-1).reshape(-1, 3)
+    k = n_all @ rec.T
+    k2 = np.sum(k * k, axis=-1)
+    keep = (k2 > 0) & (k2 <= k_max * k_max)
+    k, k2 = k[keep], k2[keep]
+    phase = k @ coord.T
+    rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge
+    s_sq = (np.abs(rho) ** 2).sum(axis=-1)
+    if kernel == "les":
+        k_hat = 4.0 * np.pi * np.exp(-k2 / (4.0 * 0.8**2)) / k2
+        k0 = 2.0 * 0.8 / np.sqrt(np.pi)
+    else:
+        amplitude = np.array([0.7, -0.2, 0.1])
+        width = np.array([0.8, 1.7, 3.2])
+        k_hat = np.sum(
+            amplitude[None, :]
+            * (2.0 * np.pi) ** 1.5
+            * width[None, :] ** 3
+            * np.exp(-0.5 * k2[:, None] * width[None, :] ** 2),
+            axis=-1,
+        )
+        k0 = amplitude.sum()
+    volume = abs(np.linalg.det(box))
+    return float(
+        COULOMB_CONSTANT
+        * ((k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum())
+    )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_fixed_grid_truncates_oversized_cells(kernel: str) -> None:
+    """A cell beyond the grid is evaluated with the grid alone, not an error."""
+    coord, atype, box = _periodic_case()
+    model = _periodic_fixed_model(kernel, box, constraint=None)
+    _set_constant_charges(model, [1.0, -0.5])
+    big_box = box * 3.0
+    with pytest.warns(UserWarning, match="truncated"):
+        out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(big_box))
+    bounds = model.atomic_model.long_range.k_grid_bounds.detach().cpu().numpy()
+    charge = np.zeros((6, CHANNELS[kernel]))
+    charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0]
+    expected = _kspace_reference_bounded(coord[0], charge, big_box[0], kernel, bounds)
+    np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-9, atol=1e-9)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_fixed_grid_needs_statistics_or_bounds(kernel: str) -> None:
+    """Without statistics or explicit bounds the grid is not sized; bounds skip them."""
+    coord, atype, box = _periodic_case()
+    unsized = _periodic_fixed_model(kernel)
+    with pytest.raises(ValueError, match="not sized yet"):
+        unsized(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box))
+    config = _config(kernel)
+    config["long_range"]["periodic"] = True
+    config["long_range"]["k_grid_mode"] = "fixed"
+    config["long_range"]["k_grid_bounds"] = [12, 12, 12]
+    model = get_model(config).to(env.DEVICE).eval()
+    np.testing.assert_array_equal(
+        model.atomic_model.long_range.k_grid_bounds.cpu().numpy(), [12, 12, 12]
+    )
+    out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box))
+    reference = _periodic_model(kernel)
+    reference.load_state_dict(model.state_dict(), strict=False)
+    expected = reference(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box))
+    torch.testing.assert_close(out["energy"], expected["energy"], rtol=1e-9, atol=1e-12)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_fixed_grid_restores_from_state_dict(kernel: str) -> None:
+    """The state dict of a sized grid loads into an unsized model, as a checkpoint does."""
+    coord, atype, box = _periodic_case()
+    model = _periodic_fixed_model(kernel, box)
+    restored = _periodic_fixed_model(kernel)
+    restored.load_state_dict(model.state_dict())
+    args = (_tensor(coord), _tensor(atype, torch.int64))
+    expected = model(*args, box=_tensor(box))
+    got = restored(*args, box=_tensor(box))
+    for key in ("energy", "force", "virial"):
+        torch.testing.assert_close(got[key], expected[key], rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_fixed_grid_serialization_round_trip(kernel: str) -> None:
+    """The record keeps the sized grid and evaluates identically."""
+    from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP
+
+    coord, atype, box = _periodic_case()
+    model = _periodic_fixed_model(kernel, box)
+    record = model.serialize()
+    assert record["long_range"]["k_grid_mode"] == "fixed"
+    assert "k_grid_bounds" in record["long_range"]
+    reference = BaseModelDP.deserialize(record)
+    assert reference.atomic_model.long_range.k_grid_mode == "fixed"
+    expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box))
+    ret = reference.call_common(coord, atype, box, neighbor_graph_method="dense")
+    np.testing.assert_allclose(
+        ret["energy_redu"].reshape(-1),
+        expected["energy"].detach().cpu().numpy().reshape(-1),
+        rtol=1e-9,
+        atol=1e-9,
+    )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+@REQUIRES_SUPPORTED_COMPILE
+def test_periodic_fixed_grid_compiled_matches_eager(kernel: str) -> None:
+    """The compiled graph lower evaluates the fixed-grid model like the eager one."""
+    coord, atype, box = _periodic_case()
+    model = _periodic_fixed_model(kernel, box)
+    coord_t, atype_t, box_t = (
+        _tensor(coord),
+        _tensor(atype, torch.int64),
+        _tensor(box),
+    )
+    model.train()
+    runner = _CompiledModel(model, _get_model_structure_key(model))
+    runner.train()
+    reference = model.eval()(coord_t, atype_t, box=box_t)
+    model.train()
+    out = runner(coord_t, atype_t, box=box_t)
+    for key in ("energy", "force", "virial"):
+        torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8)
+    loss = out["energy"].square().sum() + out["force"].square().sum()
+    loss.backward()
+    grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()]
+    assert grads and all(grad is not None for grad in grads)
+
+
+# ---------------------------------------------------------------------------
+# Periodic decoupled k grid
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_decoupled_grid_ignores_kernel_drift(kernel: str) -> None:
+    """A decoupled model selects the k points with the initial kernel but
+    evaluates them with the live one: the drift moves the energy, not the grid.
+    """
+    config = _config(kernel)
+    config["long_range"]["periodic"] = True
+    config["long_range"]["k_grid_mode"] = "ragged"
+    config["long_range"]["decouple_k_grid"] = True
+    model = get_model(config).to(env.DEVICE).eval()
+    coord, atype, box = _periodic_case()
+    coord_t, atype_t, box_t = (
+        _tensor(coord),
+        _tensor(atype, torch.int64),
+        _tensor(box),
+    )
+    long_range = model.atomic_model.long_range
+    grid_args = (
+        coord_t.reshape(-1, 3),
+        torch.tensor([6], dtype=torch.int64, device=env.DEVICE),
+        box_t,
+    )
+    build = model.atomic_model.build_topologies
+    before = build(*grid_args, atype=atype_t.reshape(-1))["k_grid"]
+    energy_before = model(coord_t, atype_t, box=box_t)["energy"]
+    # move the kernel far enough that a coupled grid would change
+    with torch.no_grad():
+        if kernel == "les":
+            long_range.kernel.log_alpha.add_(1.0)
+        else:
+            long_range.kernel.log_width.sub_(1.0)
+    after = build(*grid_args, atype=atype_t.reshape(-1))["k_grid"]
+    torch.testing.assert_close(after.vec, before.vec)
+    torch.testing.assert_close(after.mask, before.mask)
+    # the kernel values still follow the live parameters
+    energy_after = model(coord_t, atype_t, box=box_t)["energy"]
+    assert not torch.allclose(energy_after, energy_before)
+    # and the drifted kernel still takes a finite gradient through the graph
+    model.train()
+    out = model(coord_t, atype_t, box=box_t)
+    loss = out["energy"].square().sum() + out["force"].square().sum()
+    loss.backward()
+    kernel_grads = [parameter.grad for parameter in long_range.kernel.parameters()]
+    assert kernel_grads and all(
+        grad is not None and torch.isfinite(grad).all() for grad in kernel_grads
+    )
+
+
+# ---------------------------------------------------------------------------
+# Periodic PME mesh mode
+# ---------------------------------------------------------------------------
+
+
+def _periodic_pme_model(kernel: str, **kwargs: object) -> torch.nn.Module:
+    """A periodic model in the PME mesh mode, ``pme_spacing`` and ``pme_order`` in ``kwargs``."""
+    pme_kwargs = {
+        key: kwargs.pop(key) for key in ("pme_spacing", "pme_order") if key in kwargs
+    }
+    config = _config(kernel, **kwargs)
+    config["long_range"]["periodic"] = True
+    config["long_range"]["k_grid_mode"] = "pme"
+    config["long_range"].update(pme_kwargs)
+    return get_model(config).to(env.DEVICE).eval()
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_pme_matches_kspace_reference(kernel: str) -> None:
+    """With fixed charges the PME energy approaches the converged k-space sum."""
+    model = _periodic_pme_model(kernel, constraint=None)
+    _set_constant_charges(model, [1.0, -0.5])
+    coord, atype, box = _periodic_case()
+    out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box))
+    charge = np.zeros((6, CHANNELS[kernel]))
+    charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0]
+    expected = _kspace_reference(coord[0], charge, box[0], kernel)
+    # The default mesh approximates the reciprocal sum within a few parts in a
+    # million; the energy, less the self term, is several times smaller.
+    np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-5)
+    np.testing.assert_allclose(out["atom_energy"].sum().item(), expected, rtol=1e-5)
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_pme_force_virial_match_finite_differences(kernel: str) -> None:
+    """Force and strain virial are the exact derivatives of the PME energy."""
+    model = _periodic_pme_model(kernel)
+    coord, atype, box = _periodic_case()
+    coord_t, atype_t, box_t = (
+        _tensor(coord),
+        _tensor(atype, torch.int64),
+        _tensor(box),
+    )
+    out = model(coord_t, atype_t, box=box_t)
+
+    def energy(c: np.ndarray, b: np.ndarray) -> float:
+        return model(_tensor(c), atype_t, box=_tensor(b))["energy"].item()
+
+    h = 1e-5
+    force_fd = np.zeros((6, 3))
+    for i in range(6):
+        for d in range(3):
+            plus, minus = coord.copy(), coord.copy()
+            plus[0, i, d] += h
+            minus[0, i, d] -= h
+            force_fd[i, d] = -(energy(plus, box) - energy(minus, box)) / (2 * h)
+    np.testing.assert_allclose(
+        out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-5
+    )
+
+    # the strain family of the mesh strain leaf: r -> (I + eps) r and
+    # B -> B (I + eps)^T leave the fractional coordinates invariant
+    virial_fd = np.zeros((3, 3))
+    for a in range(3):
+        for b in range(3):
+            strain = np.zeros((3, 3))
+            strain[a, b] = h
+            plus = coord @ (np.eye(3) + strain).T
+            minus = coord @ (np.eye(3) - strain).T
+            box_plus = box @ (np.eye(3) + strain).T
+            box_minus = box @ (np.eye(3) - strain).T
+            virial_fd[a, b] = -(energy(plus, box_plus) - energy(minus, box_minus)) / (
+                2 * h
+            )
+    np.testing.assert_allclose(
+        out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-4
+    )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_pme_frames_are_independent(kernel: str) -> None:
+    """Frames on one mesh are independent; a larger cell sharing the batch refines it.
+
+    The frames of a batch share the mesh that resolves the largest cell, so a
+    smaller cell evaluated beside a larger one sees a finer mesh, which moves
+    its energy by less than the accuracy of the mesh.
+    """
+    model = _periodic_pme_model(kernel)
+    coord, atype, box = _periodic_case()
+    atype2 = np.concatenate([atype, atype], axis=0)
+    coord2 = np.concatenate([coord, coord + 0.1], axis=0)
+    for box2, rtol in (
+        (np.concatenate([box, box], axis=0), 1e-10),
+        (np.concatenate([box, box * 1.3], axis=0), 1e-6),
+    ):
+        batched = model(
+            _tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2)
+        )
+        for frame in range(2):
+            alone = model(
+                _tensor(coord2[frame : frame + 1]),
+                _tensor(atype2[frame : frame + 1], torch.int64),
+                box=_tensor(box2[frame : frame + 1]),
+            )
+            torch.testing.assert_close(
+                alone["energy"][0], batched["energy"][frame], rtol=rtol, atol=1e-10
+            )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_pme_serves_non_periodic_batches(kernel: str) -> None:
+    """A PME model keeps its non-periodic pair evaluation."""
+    model = _periodic_pme_model(kernel)
+    plain = _model(kernel)
+    plain.load_state_dict(model.state_dict(), strict=False)
+    coord, atype = _batch()
+    mesh_out = model(coord, atype)
+    plain_out = plain(coord, atype)
+    torch.testing.assert_close(mesh_out["energy"], plain_out["energy"])
+    torch.testing.assert_close(mesh_out["force"], plain_out["force"])
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+def test_periodic_pme_serialization_round_trip(kernel: str) -> None:
+    """The record keeps the pme mode and the order, and evaluates identically."""
+    from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP
+
+    coord, atype, box = _periodic_case()
+    model = _periodic_pme_model(kernel, pme_order=6)
+    record = model.serialize()
+    assert record["long_range"]["k_grid_mode"] == "pme"
+    assert record["long_range"]["pme_order"] == 6
+    reference = BaseModelDP.deserialize(record)
+    assert reference.atomic_model.long_range.k_grid_mode == "pme"
+    expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box))
+    ret = reference.call_common(coord, atype, box, neighbor_graph_method="dense")
+    np.testing.assert_allclose(
+        ret["energy_redu"].reshape(-1),
+        expected["energy"].detach().cpu().numpy().reshape(-1),
+        rtol=1e-9,
+        atol=1e-9,
+    )
+
+
+@pytest.mark.parametrize("kernel", ["les", "sog"])
+@REQUIRES_SUPPORTED_COMPILE
+def test_periodic_pme_compiled_matches_eager(kernel: str) -> None:
+    """The compiled graph lower evaluates the PME model like the eager one."""
+    if not mesh_compile_supported(env.DEVICE):
+        pytest.skip(f"torch {torch.__version__} miscompiles the mesh virial here")
+    model = _periodic_pme_model(kernel)
+    coord, atype, box = _periodic_case()
+    coord_t, atype_t, box_t = (
+        _tensor(coord),
+        _tensor(atype, torch.int64),
+        _tensor(box),
+    )
+    model.train()
+    runner = _CompiledModel(model, _get_model_structure_key(model))
+    runner.train()
+    reference = model.eval()(coord_t, atype_t, box=box_t)
+    model.train()
+    out = runner(coord_t, atype_t, box=box_t)
+    for key in ("energy", "force", "virial"):
+        torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8)
+    # the backward reaches the charge network and the kernel through the
+    # compiled graph
+    loss = out["energy"].square().sum() + out["force"].square().sum()
+    loss.backward()
+    grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()]
+    assert grads and all(grad is not None for grad in grads)
+    assert all(torch.isfinite(grad).all() for grad in grads)
+
+
+@pytest.mark.parametrize(
+    ("version", "lower", "topologies", "accepted"),
+    [
+        ("2.11.0+cpu", "graph", ["all_pairs", "pme_mesh"], False),
+        ("2.12.0", "graph", ["all_pairs", "pme_mesh"], True),
+        ("2.11.0+cpu", "dpa4c_canonical", ["all_pairs", "pme_mesh"], True),
+        ("2.11.0+cpu", "graph", ["all_pairs"], True),
+    ],
+)
+def test_mesh_freeze_version_guard(
+    monkeypatch, version: str, lower: str, topologies: list[str], accepted: bool
+) -> None:
+    """A .pt2 graph lower that reads the PME mesh freezes from torch 2.12 on."""
+    monkeypatch.setattr(torch, "__version__", version)
+    metadata = {"lower_input_kind": lower, "topologies": topologies}
+    if accepted:
+        check_mesh_freeze_torch_version(metadata)
+    else:
+        with pytest.raises(RuntimeError, match=r"torch >= 2\.12"):
+            check_mesh_freeze_torch_version(metadata)
+
+
+def test_periodic_freeze_checks_the_torch_version(tmp_path, monkeypatch) -> None:
+    """The .pt2 freeze of a periodic model stops before AOTInductor where it cannot run."""
+    config = _config("les")
+    config["long_range"]["periodic"] = True
+    model = get_model(config)
+    monkeypatch.setattr(serialization, "mesh_freeze_supported", lambda: False)
+    with pytest.raises(RuntimeError, match=r"torch >= 2\.12"):
+        deserialize_to_file(
+            str(tmp_path / "periodic.pt2"),
+            {"model": model.serialize(), "model_def_script": config},
+            lower_kind="graph",
+        )
+
+
+@pytest.mark.parametrize(
+    ("version", "device", "grid_mode", "accepted"),
+    [
+        ("2.12.0", "cpu", "pme", False),
+        ("2.13.0", "cpu", "pme", True),
+        ("2.11.0", "cuda", "pme", True),
+        ("2.12.0", "cpu", "ragged", True),
+    ],
+)
+def test_mesh_compile_version_guard(
+    monkeypatch, version: str, device: str, grid_mode: str, accepted: bool
+) -> None:
+    """Compiled PME training runs on the CPU from torch 2.13 on, on CUDA throughout."""
+    config = _config("les")
+    config["long_range"]["periodic"] = True
+    config["long_range"]["k_grid_mode"] = grid_mode
+    model = get_model(config)
+    monkeypatch.setattr(torch, "__version__", version)
+    if accepted:
+        check_mesh_compile_torch_version(model, torch.device(device))
+    else:
+        with pytest.raises(RuntimeError, match=r"torch >= 2\.13"):
+            check_mesh_compile_torch_version(model, torch.device(device))
+
+
+def test_periodic_pme_compile_checks_the_torch_version(monkeypatch) -> None:
+    """The compiled training trace of a PME model stops where Inductor forms a wrong virial."""
+    model = _periodic_pme_model("les")
+    monkeypatch.setattr(serialization, "mesh_compile_supported", lambda device: False)
+    with pytest.raises(RuntimeError, match=r"torch >= 2\.13"):
+        _trace_and_compile_graph(model, None, None, None, None)
diff --git a/source/tests/pt_expt/test_calculator.py b/source/tests/pt_expt/test_calculator.py
new file mode 100644
index 0000000000..147fc7b88d
--- /dev/null
+++ b/source/tests/pt_expt/test_calculator.py
@@ -0,0 +1,646 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""ASE checkpoint export and eager-inference behavior."""
+
+import contextlib
+import copy
+import gc
+import json
+import logging
+import os
+import sys
+import zipfile
+from collections.abc import (
+    Callable,
+    Iterator,
+)
+from concurrent.futures import (
+    ThreadPoolExecutor,
+)
+from functools import (
+    partial,
+)
+from pathlib import (
+    Path,
+)
+from threading import (
+    Barrier,
+)
+from types import (
+    ModuleType,
+)
+from unittest import (
+    mock,
+)
+
+import numpy as np
+import pytest
+import torch
+from ase import (
+    Atoms,
+)
+
+from deepmd import calculator as calculator_module
+from deepmd.calculator import (
+    DP,
+)
+from deepmd.infer import (
+    DeepPot,
+)
+from deepmd.pt.entrypoints import (
+    freeze_pt2,
+)
+from deepmd.pt.utils import env as pt_env
+from deepmd.pt_expt.utils import (
+    env,
+    serialization,
+)
+
+from ..common.sezm_config import (
+    _tiny_sezm_model_params,
+    _tiny_sezm_native_spin_model_params,
+)
+from .compile_utils import (
+    _COMPILE_SUPPORT_ERROR,
+    REQUIRES_SUPPORTED_COMPILE,
+)
+from .model.test_dpa4_export import (
+    _DPA4_CONFIG,
+)
+from .model.test_dpa4c_graph_lower import _config as dpa4c_config
+from .model.test_export_with_comm import (
+    _DPA3_CONFIG,
+)
+from .model.test_graph_export_with_comm import (
+    DPA1_CONFIG,
+    DPA2_CONFIG,
+)
+from .test_dp_freeze import (
+    model_se_e2_a,
+)
+
+_LEVELS = ("DP_CUDA_INFER", "DP_TRITON_INFER", "DP_CUTILE_INFER", "DP_CUTE_INFER")
+
+
+def _device_bindings() -> Iterator[tuple[ModuleType, str]]:
+    """Locate imported backend device aliases without importing more modules."""
+    for name, module in tuple(sys.modules.items()):
+        if not name.startswith(("deepmd.pt.", "deepmd.pt_expt.")) or module is None:
+            continue
+        for attribute in ("DEVICE", "device"):
+            if isinstance(vars(module).get(attribute), torch.device):
+                yield module, attribute
+
+
+@pytest.fixture(autouse=True, params=["cpu"])
+def inference_device(
+    request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch
+) -> Iterator[str]:
+    device = request.param
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    original_devices = {"pt": pt_env.DEVICE, "pt_expt": env.DEVICE}
+    selected_device = torch.device(device)
+    bindings = set(_device_bindings())
+    for module, attribute in bindings:
+        monkeypatch.setattr(module, attribute, selected_device)
+    for name in (*_LEVELS, "DP_COMPILE_INFER"):
+        monkeypatch.delenv(name, raising=False)
+    yield device
+    # Lazy imports can retain aliases that were not present at fixture setup.
+    for module, attribute in set(_device_bindings()) - bindings:
+        if getattr(module, attribute) is selected_device:
+            backend = (
+                "pt_expt" if module.__name__.startswith("deepmd.pt_expt.") else "pt"
+            )
+            setattr(module, attribute, original_devices[backend])
+
+
+def _config(kind: str) -> dict:
+    if kind in ("pt", "pt-dpa4"):
+        config = _tiny_sezm_model_params()
+        if kind == "pt-dpa4":
+            config["type"] = config["descriptor"]["type"] = "dpa4"
+    elif kind == "dpa4":
+        config = copy.deepcopy(_DPA4_CONFIG)
+        config["descriptor"].update(channels=4, n_radial=3, lmax=1, n_blocks=1)
+        config["fitting_net"]["neuron"] = [8]
+    elif kind == "dpa4c":
+        config = dpa4c_config()
+        config["descriptor"].update(channels=8, lmax=2, precision="float32")
+        config["fitting_net"]["neuron"] = [8]
+        config["fitting_net"]["precision"] = "float32"
+    else:
+        config = copy.deepcopy(
+            {
+                "dpa1": DPA1_CONFIG,
+                "dpa2": DPA2_CONFIG,
+                "dpa3": _DPA3_CONFIG,
+                "dense": model_se_e2_a,
+                "pt-dense": model_se_e2_a,
+            }[kind]
+        )
+    config["type_map"] = ["O", "H"]
+    if kind in ("dense", "pt-dense"):
+        config["descriptor"]["sel"] = config["descriptor"]["sel"][:2]
+    if env.DEVICE.type == "cuda":
+        config["descriptor"]["precision"] = "float32"
+        config["fitting_net"]["precision"] = "float32"
+        if kind in ("pt", "pt-dpa4", "dpa4"):
+            config["descriptor"]["channels"] = 32
+    return config
+
+
+def _write_checkpoint(
+    path: Path,
+    kind: str,
+    head: str | None = None,
+    *,
+    config: dict | None = None,
+) -> Path:
+    if kind.startswith("pt"):
+        from deepmd.pt.model.model import (
+            get_model,
+        )
+        from deepmd.pt.train.wrapper import (
+            ModelWrapper,
+        )
+    else:
+        from deepmd.pt_expt.model.get_model import (
+            get_model,
+        )
+        from deepmd.pt_expt.train.wrapper import (
+            ModelWrapper,
+        )
+
+    params = _config(kind) if config is None else copy.deepcopy(config)
+    model = get_model(copy.deepcopy(params)).to("cpu").eval()
+    generator = torch.Generator(device="cpu").manual_seed(17)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            if torch.count_nonzero(parameter) == 0:
+                parameter.copy_(
+                    0.05
+                    * torch.randn(
+                        parameter.shape,
+                        generator=generator,
+                        dtype=parameter.dtype,
+                        device="cpu",
+                    )
+                )
+    wrapper = ModelWrapper(
+        model if head is None else {head: model},
+        model_params=params if head is None else {"model_dict": {head: params}},
+    )
+    torch.save({"model": wrapper.state_dict()}, path)
+    return path
+
+
+@pytest.fixture
+def file_freezer(monkeypatch: pytest.MonkeyPatch) -> mock.Mock:
+    """Exercise artifact ownership without running threaded compilation."""
+
+    class FileEvaluator:
+        def __init__(self, model_file: str) -> None:
+            self.model_path = model_file
+
+        def serialize(self) -> str:
+            return Path(self.model_path).read_text()
+
+    def load(model_file: str, **kwargs: object) -> mock.Mock:
+        return mock.Mock(
+            deep_eval=FileEvaluator(model_file),
+            get_type_map=mock.Mock(return_value=["O", "H"]),
+            get_ntypes=mock.Mock(return_value=2),
+        )
+
+    def export(source: str, output: str, *, head: str | None) -> None:
+        Path(output).write_text(head or "Default")
+
+    freezer = mock.Mock(side_effect=export)
+
+    def select(
+        evaluator: object, head: str | None, **kwargs: object
+    ) -> Callable[..., None]:
+        return partial(freezer, head=head)
+
+    monkeypatch.setattr(calculator_module, "DeepPot", load)
+    monkeypatch.setattr(calculator_module, "_checkpoint_freezer", select)
+    return freezer
+
+
+def test_automatic_freeze_retains_owned_artifact(
+    tmp_path: Path, file_freezer: mock.Mock
+) -> None:
+    checkpoint = tmp_path / "model.pt"
+    first = DP(checkpoint, head="first", freeze=True)
+    second = DP(checkpoint, head="second", freeze=True)
+    evaluator = first.dp.deep_eval
+    private = Path(evaluator.model_path)
+    assert evaluator.serialize() == "first"
+    assert second.dp.deep_eval.serialize() == "second"
+    assert (tmp_path / "model-ase_single.pt2").read_text() == "second"
+    (tmp_path / "model-ase_single.pt2").write_text("external overwrite")
+    assert evaluator.serialize() == "first"
+    assert second.dp.deep_eval.serialize() == "second"
+    del first
+    gc.collect()
+    assert private.exists()
+    del evaluator
+    gc.collect()
+    assert not private.parent.exists()
+
+
+def test_concurrent_automatic_freeze_loads_each_head(
+    tmp_path: Path, file_freezer: mock.Mock
+) -> None:
+    barrier = Barrier(2)
+
+    def export(source: str, output: str, *, head: str) -> None:
+        Path(output).write_text(head)
+        barrier.wait(timeout=10)
+
+    file_freezer.side_effect = export
+
+    def construct(head: str) -> DP:
+        return DP(tmp_path / "model.pt", head=head, freeze=True)
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        first, second = pool.map(construct, ["first", "second"])
+    assert first.dp.deep_eval.serialize() == "first"
+    assert second.dp.deep_eval.serialize() == "second"
+    assert (tmp_path / "model-ase_single.pt2").read_text() in {"first", "second"}
+
+
+@pytest.mark.parametrize("stage", ["export", "load", "publish"])
+def test_failed_automatic_freeze_preserves_published_artifact(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    file_freezer: mock.Mock,
+    stage: str,
+) -> None:
+    output = tmp_path / "model-ase_single.pt2"
+    output.write_text("original")
+
+    def fail_export(source: str, output: str, **kwargs: object) -> None:
+        Path(output).write_text("incomplete")
+        raise RuntimeError("export failed")
+
+    load = calculator_module.DeepPot
+
+    def fail_load(model_file: str, **kwargs: object) -> object:
+        if model_file.endswith(".pt2"):
+            raise RuntimeError("load failed")
+        return load(model_file, **kwargs)
+
+    if stage == "export":
+        file_freezer.side_effect = fail_export
+    elif stage == "load":
+        monkeypatch.setattr(calculator_module, "DeepPot", fail_load)
+    else:
+        monkeypatch.setattr(
+            Path, "replace", mock.Mock(side_effect=RuntimeError("publish failed"))
+        )
+    with pytest.raises(RuntimeError, match=f"{stage} failed"):
+        DP(tmp_path / "model.pt", freeze=True)
+    assert output.read_text() == "original"
+    assert list(tmp_path.iterdir()) == [output]
+
+
+@pytest.mark.parametrize("skin", [None, 0.75])
+def test_graph_neighbor_options_reach_evaluator(
+    tmp_path: Path, skin: float | None
+) -> None:
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4")
+    options = {"neighbor_skin": skin} if skin is not None else {}
+    with mock.patch("deepmd.calculator.DeepPot", wraps=DeepPot) as deep_pot:
+        DP(checkpoint, neighbor_graph_method="vesin", **options)
+    deep_pot.assert_called_once_with(
+        str(checkpoint.resolve()),
+        neighbor_list=None,
+        head=None,
+        nlist_backend="auto",
+        neighbor_graph_method="vesin",
+        neighbor_skin=1.0 if skin is None else skin,
+    )
+
+
+@REQUIRES_SUPPORTED_COMPILE
+def test_native_spin_edge_freeze_preserves_neighbor_options(tmp_path: Path) -> None:
+    config = _tiny_sezm_native_spin_model_params()
+    # Edge inputs retain all cutoff neighbors independently of dense capacity.
+    config["descriptor"]["sel"] = [1, 0]
+    checkpoint = _write_checkpoint(tmp_path / "native.pt", "pt", config=config)
+    original = DeepPot(
+        str(checkpoint), neighbor_graph_method="vesin", neighbor_skin=1.0
+    )
+    frozen = tmp_path / "native.pt2"
+    freeze_pt2.freeze_sezm_to_pt2(str(checkpoint), str(frozen), with_comm=False)
+    coords = np.array([[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [-0.3, 0.8, 0.0]]])
+    cells = np.eye(3).reshape(1, 9) * 8.0
+    types = np.array([0, 1, 0])
+    spins = np.array([[[0.2, 0.3, 0.9], [0.0, 0.0, 0.0], [-0.1, 0.6, 0.7]]])
+    for method in ("vesin", "auto"):
+        exported = DeepPot(str(frozen), neighbor_graph_method=method, neighbor_skin=1.0)
+        for displacement in (0.0, 0.1, 0.7):
+            moved = coords.copy()
+            moved[0, 1, 0] += displacement
+            expected = original.eval(moved, cells, types, spin=spins, atomic=True)
+            actual = exported.eval(moved, cells, types, spin=spins, atomic=True)
+            for value, reference in zip(actual, expected, strict=True):
+                np.testing.assert_allclose(value, reference, rtol=1e-6, atol=1e-8)
+
+    native = DeepPot(str(frozen), nlist_backend="native", auto_batch_size=False)
+    batch_coords = np.repeat(coords, 2, axis=0)
+    batch_cells = np.repeat(cells, 2, axis=0)
+    batch_types = np.array([[0, 1, 0], [1, 0, 0]])
+    batch_spins = np.repeat(spins, 2, axis=0)
+    actual = native.eval(
+        batch_coords,
+        batch_cells,
+        batch_types,
+        spin=batch_spins,
+        atomic=True,
+        mixed_type=True,
+    )
+    frames = [
+        original.eval(
+            batch_coords[frame : frame + 1],
+            batch_cells[frame : frame + 1],
+            batch_types[frame],
+            spin=batch_spins[frame : frame + 1],
+            atomic=True,
+        )
+        for frame in range(2)
+    ]
+    for value, first, second in zip(actual, *frames, strict=True):
+        np.testing.assert_allclose(
+            value, np.concatenate((first, second), axis=0), rtol=1e-6, atol=1e-8
+        )
+
+
+@pytest.mark.parametrize("skin", [-1.0, float("inf"), float("nan")])
+def test_invalid_neighbor_skin_rejected(skin: float) -> None:
+    with pytest.raises(ValueError, match="neighbor_skin"):
+        DP("unused.pt", neighbor_skin=skin)
+
+
+@pytest.mark.parametrize(
+    "kind", ["pt", "pt-dpa4", "dpa4", "dpa4c", "dpa1", "dpa2", "pt-dense", "dense"]
+)
+@pytest.mark.parametrize("inference_device", ["cpu", "cuda"], indirect=True)
+def test_cached_neighbors_and_types_match_fresh_evaluation(
+    tmp_path: Path, kind: str, inference_device: str
+) -> None:
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", kind)
+    options = (
+        {"nlist_backend": "vesin"}
+        if kind in ("pt-dense", "dense")
+        else {"neighbor_graph_method": "vesin"}
+    )
+    cached = DP(checkpoint, **options)
+    fresh = DP(checkpoint, neighbor_skin=0.0, **options)
+    positions = np.array([[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [-0.3, 0.8, 0.0]])
+    for step in range(5):
+        points = positions.copy()
+        points[1, 0] += 0.1 * step
+        symbols = "OH2" if step < 2 else "HOH"
+        cell = [8.0, 8.0, 8.0] if step < 3 else [7.5, 8.0, 8.0]
+        if step == 4:
+            cached.type_dict = fresh.type_dict = {"O": 1, "H": 0}
+        atoms = Atoms(symbols, positions=points, cell=cell, pbc=True)
+        results = []
+        for calculator in (cached, fresh):
+            atoms.calc = calculator
+            results.append(
+                (atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress())
+            )
+        fp32 = kind == "dpa4c" or inference_device == "cuda"
+        rtol, atol = (1e-4, 1e-5) if fp32 else (1e-7, 1e-8)
+        for actual, expected in zip(*results, strict=True):
+            np.testing.assert_allclose(actual, expected, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("kind", ["pt", "pt-dpa4", "dpa4", "dpa4c", "dpa1", "dpa2"])
+def test_checkpoint_default_warns_without_export(
+    tmp_path: Path, caplog: pytest.LogCaptureFixture, kind: str
+) -> None:
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", kind)
+    with caplog.at_level(logging.WARNING, logger="deepmd.calculator"):
+        calculator = DP(checkpoint)
+
+    assert calculator.dp.deep_eval.model_path == str(checkpoint)
+    assert not checkpoint.with_name("model-ase_single.pt2").exists()
+    warnings = [
+        record.getMessage()
+        for record in caplog.records
+        if record.name == "deepmd.calculator"
+    ]
+    assert len(warnings) == 1
+    assert "SUBSTANTIALLY slower" in warnings[0]
+    assert "freeze=True" in warnings[0]
+    assert "manually" in warnings[0]
+    assert "debugging" in warnings[0]
+
+
+@pytest.mark.parametrize(
+    ("kind", "head"),
+    [("pt", "branch"), ("dpa4", "branch"), ("dpa4c", "branch"), ("dpa4", None)],
+)
+@pytest.mark.parametrize("override", [None, ("0", "0"), ("2", "3")])
+def test_automatic_freeze_reuses_kernel_policy(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    kind: str,
+    head: str | None,
+    override: tuple[str, str] | None,
+) -> None:
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", kind, head=head or "Default")
+    monkeypatch.setattr(freeze_pt2, "DEVICE", torch.device("cuda"))
+    export_policy = serialization._fused_operators_for_export
+
+    @contextlib.contextmanager
+    def cuda_export_policy() -> Iterator[None]:
+        # Only policy selection sees CUDA; lazy imports retain the real query.
+        with contextlib.ExitStack() as stack:
+            with mock.patch(
+                "deepmd.pt_expt.kernels.utils.backend_device_type", return_value="cuda"
+            ):
+                stack.enter_context(export_policy())
+            yield
+
+    monkeypatch.setattr(
+        serialization, "_fused_operators_for_export", cuda_export_policy
+    )
+    if override is not None:
+        for name, value in zip(_LEVELS[:2], override, strict=True):
+            monkeypatch.setenv(name, value)
+    captured = {}
+
+    def capture(*args: object, **kwargs: object) -> None:
+        captured["levels"] = tuple(os.environ.get(name) for name in _LEVELS[:2])
+        captured["with_comm"] = kwargs["with_comm"]
+        Path(args[1] if kind == "pt" else args[0]).write_bytes(b"export")
+
+    if kind == "pt":
+        monkeypatch.setattr(freeze_pt2, "_freeze_sezm_to_pt2", capture)
+    else:
+        monkeypatch.setattr(serialization, "_deserialize_to_file_pt2", capture)
+
+    def load(model_file: str, **kwargs: object) -> object:
+        if model_file.endswith(".pt"):
+            return DeepPot(model_file, **kwargs)
+        result = mock.Mock()
+        result.get_type_map.return_value = ["O", "H"]
+        result.get_ntypes.return_value = 2
+        return result
+
+    with mock.patch("deepmd.calculator.DeepPot", side_effect=load) as deep_pot:
+        DP(
+            checkpoint,
+            freeze=True,
+            head=head,
+            neighbor_graph_method="vesin",
+            neighbor_skin=0.75,
+        )
+
+    expected = override or ("1", "2")
+    if kind == "dpa4c":
+        expected = ("2", None if override is None else override[1])
+    assert captured == {"levels": expected, "with_comm": False}
+    assert tuple(os.environ.get(name) for name in _LEVELS[:2]) == (
+        override or (None, None)
+    )
+    assert deep_pot.call_count == 2
+    private = Path(deep_pot.call_args.args[0])
+    assert private.name == "model-ase_single.pt2"
+    assert private.parent != tmp_path
+    assert (tmp_path / private.name).read_bytes() == b"export"
+    deep_pot.assert_called_with(
+        str(private),
+        neighbor_list=None,
+        nlist_backend="auto",
+        neighbor_graph_method="vesin",
+        neighbor_skin=0.75,
+    )
+
+
+@pytest.mark.parametrize("kind", ["pt-dense", "dense", "dpa3"])
+@pytest.mark.parametrize("freeze", [False, True])
+def test_ineligible_checkpoint_keeps_eager_inference(
+    tmp_path: Path, caplog: pytest.LogCaptureFixture, kind: str, freeze: bool
+) -> None:
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", kind)
+    with caplog.at_level(logging.WARNING, logger="deepmd.calculator"):
+        calculator = DP(checkpoint, freeze=freeze)
+    assert calculator.dp.deep_eval.model_path == str(checkpoint)
+    assert not any(record.name == "deepmd.calculator" for record in caplog.records)
+    assert not checkpoint.with_name("model-ase_single.pt2").exists()
+
+
+def test_single_task_pt_freeze_ignores_head(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Single-task head arguments retain the eager checkpoint's semantics."""
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", "pt")
+    loaded = DeepPot(str(checkpoint), head="Default")
+
+    def write_artifact(source: str, output: str, **kwargs: object) -> None:
+        Path(output).write_bytes(b"export")
+
+    export = mock.Mock(side_effect=write_artifact)
+    monkeypatch.setattr(freeze_pt2, "freeze_sezm_to_pt2", export)
+    with mock.patch("deepmd.calculator.DeepPot", return_value=loaded):
+        DP(checkpoint, head="Default", freeze=True)
+    export.assert_called_once_with(
+        str(checkpoint),
+        mock.ANY,
+        head=None,
+        with_comm=False,
+    )
+
+
+@pytest.mark.parametrize(
+    "kind",
+    [
+        pytest.param(
+            "pt",
+            marks=pytest.mark.skipif(
+                bool(_COMPILE_SUPPORT_ERROR), reason=_COMPILE_SUPPORT_ERROR
+            ),
+        ),
+        "dpa4",
+        "dpa4c",
+        "dpa1",
+        "dpa2",
+    ],
+)
+@pytest.mark.parametrize("inference_device", ["cpu", "cuda"], indirect=True)
+def test_automatic_freeze_single_artifact_matches_eager(
+    tmp_path: Path,
+    caplog: pytest.LogCaptureFixture,
+    kind: str,
+    inference_device: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    config = _config(kind)
+    options = {}
+    if kind == "pt":
+        config["descriptor"]["sel"] = [1, 0]
+        config["pair_exclude_types"] = [[0, 1]]
+        options["nlist_backend"] = "native"
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", kind, config=config)
+    original = checkpoint.read_bytes()
+    atoms = Atoms(
+        "OH2",
+        positions=[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [-0.3, 0.8, 0.0]],
+        cell=[8.0, 8.0, 8.0],
+        pbc=True,
+    )
+    atoms.calc = DP(checkpoint, **options)
+    expected = (atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress())
+    if kind == "pt":
+        np.testing.assert_array_equal(expected[1][0], 0.0)
+        from deepmd.pt_expt.utils import (
+            auto_neighbor_list,
+        )
+
+        monkeypatch.setattr(auto_neighbor_list, "is_nv_available", lambda: False)
+        monkeypatch.setattr(
+            auto_neighbor_list, "is_vesin_torch_available", lambda: False
+        )
+    caplog.clear()
+    with mock.patch(
+        "torch._inductor.aoti_compile_and_package",
+        wraps=torch._inductor.aoti_compile_and_package,
+    ) as compile_package:
+        atoms.calc = DP(checkpoint, freeze=True, **options)
+    assert compile_package.call_count == 1
+    if kind == "dpa4c":
+        exported = compile_package.call_args.args[0]
+        targets = {str(node.target) for node in exported.graph_module.graph.nodes}
+        assert any("dpa4c_" in target and "compress" in target for target in targets)
+    frozen = tmp_path / "model-ase_single.pt2"
+    assert checkpoint.read_bytes() == original
+    private = Path(atoms.calc.dp.deep_eval.model_path)
+    assert private != frozen
+    assert private.read_bytes() == frozen.read_bytes()
+    with zipfile.ZipFile(frozen) as archive:
+        metadata = json.loads(archive.read("model/extra/metadata.json"))
+        assert metadata["has_comm_artifact"] is False
+        assert "model/extra/forward_lower_with_comm.pt2" not in archive.namelist()
+    actual = (atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress())
+    assert np.max(np.abs(expected[1])) > 1e-8
+    fp32 = inference_device == "cuda" or kind == "dpa4c"
+    rtol, atol = (1e-4, 1e-5) if fp32 else (1e-6, 1e-8)
+    for reference, result in zip(expected, actual, strict=True):
+        np.testing.assert_allclose(result, reference, rtol=rtol, atol=atol)
+    assert not any(
+        record.name == "deepmd.calculator" and record.levelno >= logging.WARNING
+        for record in caplog.records
+    )
+    atoms.calc = DP(frozen, freeze=True)
+    reload_tol = 1e-6 if fp32 else 1e-12
+    np.testing.assert_allclose(
+        atoms.get_forces(), actual[1], rtol=reload_tol, atol=reload_tol
+    )
+    assert not (tmp_path / "model-ase_single-ase_single.pt2").exists()
diff --git a/source/tests/pt_expt/test_change_bias.py b/source/tests/pt_expt/test_change_bias.py
index 2401aa59c5..f1212324bf 100644
--- a/source/tests/pt_expt/test_change_bias.py
+++ b/source/tests/pt_expt/test_change_bias.py
@@ -19,18 +19,27 @@
 from deepmd.pt_expt.model.get_model import (
     get_model,
 )
+from deepmd.pt_expt.train.training import (
+    model_change_out_bias,
+)
 from deepmd.pt_expt.train.wrapper import (
     ModelWrapper,
 )
 from deepmd.pt_expt.utils.env import (
     DEVICE,
 )
+from deepmd.pt_expt.utils.stat import (
+    make_stat_input,
+)
 from deepmd.utils.argcheck import (
     normalize,
 )
 from deepmd.utils.compat import (
     update_deepmd_input,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 
 def to_numpy(x):
@@ -184,6 +193,7 @@ def setUpClass(cls) -> None:
 
         # Build & train 1-step model
         config = _make_config(data_dir)
+        config["model"]["preset_out_bias"] = {"energy": [-10.0, -13.6]}
         config = update_deepmd_input(config, warning=False)
         config = normalize(config)
         config["training"]["save_ckpt"] = "model.ckpt"
@@ -193,8 +203,18 @@ def setUpClass(cls) -> None:
 
         cls.model_path = os.path.join(cls.tmpdir, "model.ckpt.pt")
 
-        # Record original bias
-        cls.original_bias = to_numpy(trainer.wrapper.model["Default"].get_out_bias())
+        # A checkpoint may already carry a downstream-adjusted finite bias.
+        model = trainer.wrapper.model["Default"]
+        old_bias = model.get_out_bias()
+        model.set_out_bias(
+            torch.as_tensor(
+                [7.0, 8.0], dtype=old_bias.dtype, device=old_bias.device
+            ).reshape(old_bias.shape)
+        )
+        checkpoint = torch.load(cls.model_path, map_location=DEVICE, weights_only=True)
+        checkpoint["model"] = trainer.wrapper.state_dict()
+        torch.save(checkpoint, cls.model_path)
+        cls.original_bias = to_numpy(model.get_out_bias()).copy()
 
         # Pre-freeze shared .pte and .pt2 files so individual tests don't
         # each pay the AOTInductor compilation cost (~82s per .pt2).
@@ -231,6 +251,44 @@ def _load_model_from_ckpt(self, ckpt_path: str):
         wrapper.load_state_dict(model_state)
         return model
 
+    def _expected_fitted_bias(self):
+        model = self._load_model_from_ckpt(self.model_path)
+        # Compare against the established no-preset statistics path with the
+        # same network and adjusted bias, rather than starting from the preset.
+        model.atomic_model.preset_out_bias = None
+        data = DeepmdDataSystem(
+            systems=self.data_file,
+            batch_size=1,
+            test_size=1,
+            rcut=model.get_rcut(),
+            type_map=model.get_type_map(),
+        )
+        from deepmd.pt_expt.train.training import (
+            get_loss,
+        )
+
+        loss = get_loss({"inference": True}, 1.0, len(model.get_type_map()), model)
+        data.add_data_requirements(loss.label_requirement)
+        sampled = make_stat_input(data, min(data.get_nbatches()))
+        return to_numpy(model_change_out_bias(model, sampled).get_out_bias())
+
+    def _assert_preset_bias_update(self, model) -> None:
+        np.testing.assert_allclose(
+            to_numpy(model.get_out_bias()), self._expected_fitted_bias(), atol=1e-10
+        )
+        self.assertEqual(
+            model.atomic_model.preset_out_bias, {"energy": [[-10.0], [-13.6]]}
+        )
+
+    def test_default_bias_helper_preserves_preset_behavior(self) -> None:
+        model = self._load_model_from_ckpt(self.model_path)
+        np.testing.assert_allclose(to_numpy(model.get_out_bias()), self.original_bias)
+        sampled = [{"atype": np.array([[0, 1]], dtype=np.int32)}]
+        model_change_out_bias(model, sampled)
+        np.testing.assert_allclose(
+            to_numpy(model.get_out_bias()).reshape(-1), [-10.0, -13.6]
+        )
+
     def test_change_bias_with_data(self) -> None:
         output_path = os.path.join(self.tmpdir, "model_data_bias.pt")
         run_dp(
@@ -239,6 +297,7 @@ def test_change_bias_with_data(self) -> None:
         )
         updated_model = self._load_model_from_ckpt(output_path)
         updated_bias = to_numpy(updated_model.get_out_bias())
+        self._assert_preset_bias_update(updated_model)
         original_bias = np.array(self.original_bias)
         # Bias should have changed from the original
         self.assertFalse(
@@ -333,6 +392,11 @@ def test_change_bias_frozen_pte(self) -> None:
         updated_data = serialize_from_file(output_pte)
         updated_model = BaseModel.deserialize(updated_data["model"])
         updated_bias = to_numpy(updated_model.get_out_bias())
+        self._assert_preset_bias_update(updated_model)
+        self.assertEqual(
+            updated_data["model_def_script"]["preset_out_bias"],
+            self.config["model"]["preset_out_bias"],
+        )
 
         # Bias should have changed
         self.assertFalse(
@@ -362,6 +426,11 @@ def test_change_bias_frozen_pt2(self) -> None:
         updated_data = serialize_from_file(output_pt2)
         updated_model = BaseModel.deserialize(updated_data["model"])
         updated_bias = to_numpy(updated_model.get_out_bias())
+        self._assert_preset_bias_update(updated_model)
+        self.assertEqual(
+            updated_data["model_def_script"]["preset_out_bias"],
+            self.config["model"]["preset_out_bias"],
+        )
 
         self.assertFalse(
             np.allclose(original_bias, updated_bias),
@@ -435,6 +504,29 @@ def test_change_bias_pte_preserves_model_def_script(self) -> None:
         self.assertIn("model_def_script", updated_data)
         self.assertEqual(updated_data["model_def_script"], original_mds)
 
+    def test_change_bias_keeps_the_per_atom_virial(self) -> None:
+        """change_bias keeps the per-atom virial of a dense-lower archive."""
+        from deepmd.pt_expt.utils.serialization import (
+            deserialize_to_file,
+            serialize_from_file,
+        )
+
+        data = serialize_from_file(self.shared_pte)
+        self.assertEqual(data["lower_input_kind"], "nlist")
+        self.assertFalse(data["do_atomic_virial"])
+        with_virial = os.path.join(self.tmpdir, "frozen_atomic_virial.pte")
+        deserialize_to_file(
+            with_virial,
+            {"model": data["model"], "model_def_script": data["model_def_script"]},
+            do_atomic_virial=True,
+            lower_kind="nlist",
+        )
+        output_pte = os.path.join(self.tmpdir, "frozen_atomic_virial_updated.pte")
+        run_dp(f"dp --pt-expt change-bias {with_virial} -b 0.1 3.2 -o {output_pte}")
+        updated = serialize_from_file(output_pte)
+        self.assertEqual(updated["lower_input_kind"], "nlist")
+        self.assertTrue(updated["do_atomic_virial"])
+
     def test_change_bias_pt2_preserves_model_def_script(self) -> None:
         """change_bias on .pt2 should preserve model_def_script (training config)."""
         from deepmd.pt_expt.utils.serialization import (
@@ -472,6 +564,8 @@ def _make_mock_model(self):
         fitting_net = MagicMock()
 
         class FakeModel(DPModelCommon):
+            atomic_model = MagicMock(**{"has_long_range.return_value": False})
+
             def get_out_bias(self):
                 return np.array([[0.0, 0.0]])
 
diff --git a/source/tests/pt_expt/test_checkpoint_guard.py b/source/tests/pt_expt/test_checkpoint_guard.py
new file mode 100644
index 0000000000..54525725ae
--- /dev/null
+++ b/source/tests/pt_expt/test_checkpoint_guard.py
@@ -0,0 +1,442 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Checkpoint preflights exercise real trainer writers without a training dataset."""
+
+from copy import (
+    deepcopy,
+)
+from datetime import (
+    timedelta,
+)
+from functools import (
+    partial,
+)
+from pathlib import (
+    Path,
+)
+from types import (
+    SimpleNamespace,
+)
+from unittest.mock import (
+    patch,
+)
+
+import numpy as np
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.distributed.optim import (
+    ZeroRedundancyOptimizer,
+)
+from torch.nn.parallel import (
+    DistributedDataParallel,
+)
+
+from deepmd.dpmodel.train import (
+    ShardingPolicy,
+)
+from deepmd.dpmodel.train.checkpoint import (
+    CheckpointStore,
+)
+from deepmd.pt.train.training import Trainer as PTTrainer
+from deepmd.pt.train.wrapper import ModelWrapper as PTModelWrapper
+from deepmd.pt_expt.train.ema import (
+    EMA_CHECKPOINT_KEY,
+    EMA_VALIDATION_STATE_KEY,
+    ModelEMA,
+)
+from deepmd.pt_expt.train.gradient import (
+    NonFiniteGradGuard,
+)
+from deepmd.pt_expt.train.training import Trainer as PTExptTrainer
+from deepmd.pt_expt.train.validation import (
+    TOPK_RECORDS_INFO_KEY,
+    build_full_validators,
+)
+from deepmd.pt_expt.train.wrapper import ModelWrapper as PTExptModelWrapper
+
+
+def _make_trainer(backend, directory, *, rank=0, stage=0, distributed=False):
+    """Initialize only checkpoint dependencies, retaining production methods."""
+    trainer_class, wrapper_class = (
+        (PTTrainer, PTModelWrapper)
+        if backend == "pt"
+        else (PTExptTrainer, PTExptModelWrapper)
+    )
+    trainer = trainer_class.__new__(trainer_class)
+    trainer.rank = rank
+    trainer.is_distributed = distributed
+    if backend == "pt_expt":
+        trainer.training_data_by_task = {}
+        trainer._configured_data_seed = None
+        trainer._task_rng = np.random.default_rng(0)
+        trainer.world_size = dist.get_world_size() if distributed else 1
+    trainer.sharding = ShardingPolicy(stage=stage)
+    trainer.model = torch.nn.Linear(2, 1, device="cpu")
+    with torch.no_grad():
+        trainer.model.weight.fill_(2.0)
+        trainer.model.bias.fill_(2.0)
+    trainer.model_ema = ModelEMA(trainer.model, decay=0.9)
+    with torch.no_grad():
+        trainer.model.weight.fill_(3.0)
+        trainer.model.bias.fill_(3.0)
+    wrapper = wrapper_class(trainer.model)
+    trainer.wrapper = DistributedDataParallel(wrapper) if distributed else wrapper
+    trainer.optimizer = (
+        ZeroRedundancyOptimizer(
+            trainer.wrapper.parameters(), optimizer_class=torch.optim.Adam, lr=1e-3
+        )
+        if stage == 1
+        else torch.optim.Adam(trainer.wrapper.parameters(), lr=1e-3)
+    )
+    trainer.nonfinite_grad_guard = NonFiniteGradGuard()
+    trainer.ckpt_store = CheckpointStore(directory / "model.ckpt")
+    trainer.ema_ckpt_store = CheckpointStore(directory / "model_ema.ckpt")
+    directory.mkdir(parents=True, exist_ok=True)
+    return trainer
+
+
+def _make_validators(trainer, directory, *, save_best=True, restart=False):
+    wrapper = trainer.wrapper.module if trainer.is_distributed else trainer.wrapper
+    return build_full_validators(
+        validating_params={
+            "full_validation": True,
+            "ema_full_validation": True,
+            "validation_freq": 1,
+            "full_val_start": 0.0,
+            "save_best": save_best,
+            "max_best_ckpt": 1,
+            "validation_metric": "E:MAE",
+            "full_val_file": str(directory / "val.log"),
+        },
+        validation_data=SimpleNamespace(systems=[]),
+        model=trainer.model,
+        state_store=wrapper.train_infos,
+        num_steps=10,
+        rank=trainer.rank,
+        restart_training=restart,
+        checkpoint_dir=directory,
+        ensure_supported=lambda: None,
+        model_ema=trainer.model_ema,
+        sharding=trainer.sharding,
+    )
+
+
+def _best_writer(trainer, ema):
+    if isinstance(trainer, PTTrainer):
+        return trainer.save_ema_model if ema else trainer.save_model
+    return (
+        trainer._save_full_validation_ema_checkpoint
+        if ema
+        else trainer._save_full_validation_checkpoint
+    )
+
+
+def _validate(trainer, validator, *, step, metric, ema=False):
+    save_checkpoint = _best_writer(trainer, ema)
+    if isinstance(trainer, PTExptTrainer) and not ema:
+        save_checkpoint = partial(
+            save_checkpoint, data_progress=trainer._collect_data_progress()
+        )
+    with patch.object(
+        validator, "evaluate_all_systems", return_value={validator.metric_key: metric}
+    ):
+        return validator.run(
+            step_id=step,
+            display_step=step,
+            lr=1e-3,
+            save_checkpoint=save_checkpoint,
+            prepare_checkpoint=trainer.ensure_finite_gradients_for_checkpoint,
+        )
+
+
+def _record_recovered_gradients(trainer):
+    """A later finite gradient must not hide the earlier failed training step."""
+    trainer.nonfinite_grad_guard.update(torch.tensor(float("nan"), device="cpu"))
+    for parameter in trainer.wrapper.parameters():
+        parameter.grad = torch.ones_like(parameter)
+    trainer.nonfinite_grad_guard.update(torch.tensor(1.0, device="cpu"))
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+@pytest.mark.parametrize("ema", [False, True])
+def test_best_checkpoint_rejects_prior_nan_without_publication(backend, ema, tmp_path):
+    trainer = _make_trainer(backend, tmp_path)
+    validator = _make_validators(trainer, tmp_path)[int(ema)]
+    previous = _validate(trainer, validator, step=1, metric=2.0, ema=ema)
+    old_path = Path(previous.saved_best_path)
+    old_bytes = old_path.read_bytes()
+    old_records = list(validator.topk_records)
+    old_state = deepcopy(validator.state_store)
+    _record_recovered_gradients(trainer)
+
+    with (
+        patch.object(
+            trainer,
+            "_collect_checkpoint_states",
+            wraps=trainer._collect_checkpoint_states,
+        ) as collect,
+        patch("torch.save", wraps=torch.save) as serialize,
+        pytest.raises(RuntimeError, match="Non-finite gradient norm"),
+    ):
+        _validate(trainer, validator, step=2, metric=1.0, ema=ema)
+
+    collect.assert_not_called()
+    serialize.assert_not_called()
+    assert validator.topk_records == old_records
+    assert validator.state_store == old_state
+    assert old_path.read_bytes() == old_bytes
+    assert list(tmp_path.glob(f"{validator.best_checkpoint_prefix}-*.pt")) == [old_path]
+
+
+@pytest.mark.parametrize("ema", [False, True])
+def test_pt_expt_run_full_validation_wires_checkpoint_preflight(ema, tmp_path):
+    trainer = _make_trainer("pt_expt", tmp_path)
+    live_validator, ema_validator = _make_validators(trainer, tmp_path)
+    trainer.full_validator = None if ema else live_validator
+    trainer.ema_full_validator = ema_validator if ema else None
+    validator = ema_validator if ema else live_validator
+    previous_state = deepcopy(validator.state_store)
+    _record_recovered_gradients(trainer)
+    with (
+        patch.object(
+            validator,
+            "evaluate_all_systems",
+            return_value={validator.metric_key: 1.0},
+        ),
+        patch("torch.save", wraps=torch.save) as serialize,
+        pytest.raises(RuntimeError, match="Non-finite gradient norm"),
+    ):
+        trainer.run_full_validation(step=0, display_step=1, learning_rate=1e-3)
+    serialize.assert_not_called()
+    assert validator.state_store == previous_state
+    assert not validator.topk_records
+    assert not list(tmp_path.glob("*.pt"))
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+def test_finite_best_checkpoints_serialize_topk_and_restart(backend, tmp_path):
+    trainer = _make_trainer(backend, tmp_path)
+    validators = _make_validators(trainer, tmp_path)
+    for step, metric in ((1, 2.0), (2, 1.0)):
+        for ema, validator in enumerate(validators):
+            trainer.nonfinite_grad_guard.update(torch.tensor(1.0, device="cpu"))
+            result = _validate(
+                trainer, validator, step=step, metric=metric, ema=bool(ema)
+            )
+            checkpoint = torch.load(result.saved_best_path, weights_only=True)
+            expected_weight = 2.0 if ema else 3.0
+            assert torch.all(
+                checkpoint["model"]["model.Default.weight"] == expected_weight
+            )
+            assert torch.all(trainer.model.weight == 3.0)
+            if not ema:
+                assert checkpoint["model"]["_extra_state"]["train_infos"][
+                    TOPK_RECORDS_INFO_KEY
+                ] == [{"metric": metric, "step": step}]
+            else:
+                assert "optimizer" not in checkpoint
+                assert EMA_CHECKPOINT_KEY not in checkpoint
+
+    # A regular live checkpoint carries both validators' restart bookkeeping;
+    # EMA-only deployment snapshots deliberately omit EMA/optimizer state.
+    restart_path = tmp_path / "restart.pt"
+    _best_writer(trainer, False)(restart_path, step=2)
+    checkpoint = torch.load(restart_path, weights_only=True)
+    expected_records = [{"metric": 1.0, "step": 2}]
+    assert (
+        checkpoint[EMA_CHECKPOINT_KEY][EMA_VALIDATION_STATE_KEY][TOPK_RECORDS_INFO_KEY]
+        == expected_records
+    )
+    restarted = _make_trainer(backend, tmp_path)
+    restarted.wrapper.load_state_dict(checkpoint["model"])
+    restarted.model_ema.load_state_dict(checkpoint[EMA_CHECKPOINT_KEY])
+    for validator in _make_validators(restarted, tmp_path, restart=True):
+        assert validator.state_store[TOPK_RECORDS_INFO_KEY] == expected_records
+        assert [(record.metric, record.step) for record in validator.topk_records] == [
+            (1.0, 2)
+        ]
+        paths = list(tmp_path.glob(f"{validator.best_checkpoint_prefix}-*.pt"))
+        assert len(paths) == 1
+        assert paths[0].name == f"{validator.best_checkpoint_prefix}-2.t-1.pt"
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+@pytest.mark.parametrize("ema", [False, True])
+@pytest.mark.parametrize(
+    "no_save_reason", ["not_selected", "save_best_false", "not_due"]
+)
+def test_validation_without_save_keeps_accumulated_failure(
+    backend, ema, no_save_reason, tmp_path
+):
+    trainer = _make_trainer(backend, tmp_path)
+    validator = _make_validators(
+        trainer, tmp_path, save_best=no_save_reason != "save_best_false"
+    )[int(ema)]
+    if no_save_reason == "not_selected":
+        _validate(trainer, validator, step=1, metric=1.0, ema=ema)
+    elif no_save_reason == "not_due":
+        validator.validation_freq = 3
+    _record_recovered_gradients(trainer)
+    flag = trainer.nonfinite_grad_guard._nonfinite
+    with (
+        patch.object(
+            trainer,
+            "ensure_finite_gradients_for_checkpoint",
+            wraps=trainer.ensure_finite_gradients_for_checkpoint,
+        ) as prepare,
+        patch("torch.save", wraps=torch.save) as serialize,
+    ):
+        result = _validate(trainer, validator, step=2, metric=2.0, ema=ema)
+    assert result is None or result.saved_best_path is None
+    prepare.assert_not_called()
+    serialize.assert_not_called()
+    assert trainer.nonfinite_grad_guard._nonfinite is flag
+    with pytest.raises(RuntimeError, match="Non-finite gradient norm"):
+        trainer.ensure_finite_gradients_for_checkpoint()
+
+
+def _save_regular_checkpoint(trainer, step):
+    if isinstance(trainer, PTTrainer):
+        trainer._save_training_checkpoint(step=step, display_step=step, lr=1e-3)
+    else:
+        trainer.save_checkpoint(step)
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+def test_regular_checkpoint_checks_once_for_live_and_ema(backend, tmp_path):
+    trainer = _make_trainer(backend, tmp_path)
+    trainer.nonfinite_grad_guard.update(torch.tensor(1.0, device="cpu"))
+    with patch.object(
+        trainer.nonfinite_grad_guard,
+        "raise_if_nonfinite",
+        wraps=trainer.nonfinite_grad_guard.raise_if_nonfinite,
+    ) as check:
+        _save_regular_checkpoint(trainer, 1)
+    check.assert_called_once()
+    assert trainer.ckpt_store.path_for(1).is_file()
+    assert trainer.ema_ckpt_store.path_for(1).is_file()
+    _record_recovered_gradients(trainer)
+    with (
+        patch.object(trainer, "_collect_checkpoint_states") as collect,
+        pytest.raises(RuntimeError, match="Non-finite gradient norm"),
+    ):
+        _save_regular_checkpoint(trainer, 2)
+    collect.assert_not_called()
+    assert not trainer.ckpt_store.path_for(2).exists()
+    assert not trainer.ema_ckpt_store.path_for(2).exists()
+
+
+def _distributed_checkpoint_worker(rank, rendezvous, directory):
+    """Run real Gloo and ZeRO-1 collection after a rank-local sticky failure."""
+    torch.set_num_threads(1)
+    dist.init_process_group(
+        "gloo",
+        init_method=f"file://{rendezvous}",
+        rank=rank,
+        world_size=2,
+        timeout=timedelta(seconds=20),
+    )
+    try:
+        for backend in ("pt", "pt_expt"):
+            for stage in (0, 1):
+                for boundary in ("live_best", "ema_best", "regular"):
+                    path = Path(directory) / f"{backend}-{stage}-{boundary}"
+                    trainer = _make_trainer(
+                        backend, path, rank=rank, stage=stage, distributed=True
+                    )
+                    if backend == "pt_expt":
+                        assert trainer.checkpoint_is_collective
+                    ema = boundary == "ema_best"
+                    validator = _make_validators(trainer, path)[int(ema)]
+                    old_state = deepcopy(validator.state_store)
+
+                    def save(step):
+                        if boundary == "regular":
+                            _save_regular_checkpoint(trainer, step)
+                        else:
+                            _validate(
+                                trainer, validator, step=step, metric=0.5, ema=ema
+                            )
+
+                    # Rank 0 has no accumulated norm at all. It must still join
+                    # the preflight and see rank 1's historical failure before
+                    # any checkpoint state collection (in particular ZeRO).
+                    if rank == 1:
+                        _record_recovered_gradients(trainer)
+                    with (
+                        patch.object(
+                            trainer,
+                            "_collect_checkpoint_states",
+                            wraps=trainer._collect_checkpoint_states,
+                        ) as collect,
+                        patch("torch.save", wraps=torch.save) as serialize,
+                        pytest.raises(RuntimeError, match="Non-finite gradient norm"),
+                    ):
+                        save(1)
+                    collect.assert_not_called()
+                    serialize.assert_not_called()
+                    assert validator.state_store == old_state
+                    assert not validator.topk_records
+                    assert not list(path.glob("*.pt"))
+
+                    # A finite interval still saves normally. Materialize Adam
+                    # state so live saves exercise actual ZeRO consolidation.
+                    for parameter in trainer.wrapper.parameters():
+                        parameter.grad = torch.ones_like(parameter)
+                    trainer.optimizer.step()
+                    trainer.nonfinite_grad_guard.update(torch.tensor(1.0, device="cpu"))
+                    with patch.object(
+                        trainer,
+                        "_collect_checkpoint_states",
+                        wraps=trainer._collect_checkpoint_states,
+                    ) as collect:
+                        save(2)
+                    expected_calls = 2 if boundary == "regular" else 1
+                    if stage == 0 and rank != 0:
+                        expected_calls = 0
+                    assert collect.call_count == expected_calls
+                    if rank == 0:
+                        checkpoint_path = (
+                            trainer.ckpt_store.path_for(2)
+                            if boundary == "regular"
+                            else path / f"{validator.best_checkpoint_prefix}-2.t-1.pt"
+                        )
+                        checkpoint = torch.load(checkpoint_path, weights_only=True)
+                        if ema:
+                            assert "optimizer" not in checkpoint
+                        else:
+                            assert len(checkpoint["optimizer"]["state"]) == 2
+                        if boundary != "regular":
+                            assert validator.state_store[TOPK_RECORDS_INFO_KEY] == [
+                                {"metric": 0.5, "step": 2}
+                            ]
+                    dist.barrier()
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(
+    not dist.is_available() or not dist.is_gloo_available(), reason="Gloo is required"
+)
+def test_rank_local_failure_precedes_zero_checkpoint_collectives(tmp_path):
+    context = mp.spawn(
+        _distributed_checkpoint_worker,
+        args=(str(tmp_path / "rendezvous"), str(tmp_path)),
+        nprocs=2,
+        join=False,
+    )
+    try:
+        # A mismatched collective fails its process-group timeout; this outer
+        # deadline also prevents a regression from wedging the test runner.
+        if not context.join(timeout=90):
+            if not context.join(timeout=30):
+                pytest.fail("Checkpoint workers did not finish within 120 seconds")
+    finally:
+        for process in context.processes:
+            if process.is_alive():
+                process.terminate()
+            process.join(timeout=10)
+            if process.is_alive():
+                process.kill()
+                process.join(timeout=10)
diff --git a/source/tests/pt_expt/test_data_progress_restart.py b/source/tests/pt_expt/test_data_progress_restart.py
new file mode 100644
index 0000000000..9ad74a6385
--- /dev/null
+++ b/source/tests/pt_expt/test_data_progress_restart.py
@@ -0,0 +1,1023 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Training-stream continuity through native pt_expt checkpoints."""
+
+from copy import (
+    deepcopy,
+)
+from pathlib import (
+    Path,
+)
+from shutil import (
+    copyfile,
+    copytree,
+)
+from typing import (
+    Any,
+)
+from unittest.mock import (
+    Mock,
+    patch,
+)
+
+import h5py
+import lmdb
+import msgpack
+import numpy as np
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from deepmd.dpmodel.train import (
+    RankContext,
+)
+from deepmd.dpmodel.train.data_progress import (
+    DATA_PROGRESS_CHECKPOINT_KEY,
+    collect_training_data_progress,
+    preserve_training_data_progress,
+)
+from deepmd.pt_expt.entrypoints.main import (
+    get_trainer,
+)
+from deepmd.pt_expt.train import training as training_module
+from deepmd.pt_expt.train.ema import (
+    EMA_CHECKPOINT_KEY,
+)
+from deepmd.pt_expt.train.training import (
+    Trainer,
+)
+from deepmd.pt_expt.utils.env import (
+    DEVICE,
+)
+from deepmd.pt_expt.utils.lmdb_dataset import (
+    LmdbDataSystem,
+)
+from deepmd.utils import random as dp_random
+from deepmd.utils.argcheck import (
+    normalize,
+)
+from deepmd.utils.compat import (
+    update_deepmd_input,
+)
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
+
+from ..common.dpmodel.test_lmdb_data import (
+    _create_lmdb,
+)
+
+
+def _write_directory(path: Path) -> None:
+    path.mkdir()
+    np.savetxt(path / "type.raw", [0, 1, 1], fmt="%d")
+    (path / "type_map.raw").write_text("O H\n")
+    rng = np.random.RandomState(24)
+    for index in range(2):
+        target = path / f"set.{index:03d}"
+        target.mkdir()
+        np.save(target / "coord.npy", rng.uniform(0.5, 2.5, (2, 9)))
+        np.save(target / "box.npy", np.tile(np.eye(3).reshape(1, 9) * 8, (2, 1)))
+        np.save(target / "energy.npy", np.arange(index * 2, index * 2 + 2) * 0.01)
+        np.save(target / "force.npy", rng.normal(size=(2, 9)) * 0.01)
+
+
+def _write_lmdb(path: Path, *, mixed: bool) -> None:
+    def encode(array: np.ndarray) -> dict:
+        return {
+            "nd": None,
+            "type": str(array.dtype),
+            "kind": "",
+            "shape": list(array.shape),
+            "data": array.tobytes(),
+        }
+
+    nlocs = [3, 4, 3, 4, 3, 4] if mixed else [3] * 4
+    rng = np.random.RandomState(24)
+    with lmdb.open(str(path), map_size=10 * 1024 * 1024) as env:
+        with env.begin(write=True) as txn:
+            txn.put(
+                b"__metadata__",
+                msgpack.packb(
+                    {
+                        "nframes": len(nlocs),
+                        "frame_idx_fmt": "012d",
+                        "frame_nlocs": nlocs,
+                        "type_map": ["O", "H"],
+                        "system_info": {"natoms": [1, 2], "nframes": len(nlocs)},
+                    },
+                    use_bin_type=True,
+                ),
+            )
+            for index, nloc in enumerate(nlocs):
+                frame = {
+                    "atom_numbs": [1, nloc - 1],
+                    "atom_names": ["O", "H"],
+                    "atom_types": encode(np.array([0] + [1] * (nloc - 1))),
+                    "orig": encode(np.zeros(3)),
+                    "cells": encode(np.eye(3) * 8),
+                    "coords": encode(rng.uniform(0.5, 2.5, (nloc, 3))),
+                    "energies": encode(np.array(index * 0.01)),
+                    "forces": encode(rng.normal(size=(nloc, 3)) * 0.01),
+                }
+                txn.put(
+                    f"{index:012d}".encode(),
+                    msgpack.packb(frame, use_bin_type=True),
+                )
+
+
+def _write_hdf5(path: Path) -> None:
+    """Store the directory fixture in a single HDF5 system."""
+    directory = path.with_suffix("")
+    _write_directory(directory)
+    with h5py.File(path, "w") as handle:
+        handle.create_dataset("system/type.raw", data=[0, 1, 1])
+        handle.create_dataset("system/type_map.raw", data=[b"O", b"H"])
+        for array in directory.glob("set.*/*.npy"):
+            handle.create_dataset(
+                "system/" + array.relative_to(directory).as_posix(), data=np.load(array)
+            )
+
+
+def _config(path: str | Path, *, mode: str, multitask: bool = True) -> dict:
+    descriptor = {
+        "type": "se_e2_a",
+        "sel": [4, 4],
+        "rcut_smth": 0.5,
+        "rcut": 3.0,
+        "neuron": [4, 8],
+        "axis_neuron": 2,
+        "seed": 1,
+    }
+    if mode == "ragged":
+        descriptor.update(type="dpa1", sel=8, attn_layer=0)
+    model = {
+        "type_map": ["O", "H"],
+        "descriptor": descriptor,
+        "fitting_net": {"neuron": [4], "seed": 1},
+        "data_stat_nbatch": 1,
+    }
+    loss = {
+        "type": "ener",
+        "start_pref_e": 1,
+        "limit_pref_e": 1,
+        "start_pref_f": 1,
+        "limit_pref_f": 1,
+        "start_pref_v": 0,
+        "limit_pref_v": 0,
+    }
+    dataset = {
+        "systems": [str(path)] if mode == "directory" else str(path),
+        "batch_size": "mix:7" if mode == "ragged" else 1,
+    }
+    task_data = {
+        "training_data": deepcopy(dataset),
+        "validation_data": {**dataset, "numb_btch": 2},
+    }
+    training = {
+        "numb_steps": 20,
+        "seed": 19,
+        "disp_freq": 3,
+        "save_freq": 10,
+    }
+    config = {
+        "learning_rate": {
+            "type": "exp",
+            "start_lr": 1e-3,
+            "stop_lr": 1e-5,
+            "decay_steps": 5,
+        },
+        "training": training,
+    }
+    if multitask:
+        keys = ["first", "second"]
+        config["model"] = {"model_dict": {key: deepcopy(model) for key in keys}}
+        config["loss_dict"] = {key: deepcopy(loss) for key in keys}
+        training["data_dict"] = {key: deepcopy(task_data) for key in keys}
+        training["model_prob"] = {"first": 0.5, "second": 0.5}
+    else:
+        config.update(model=model, loss=loss)
+        training.update(task_data)
+    return normalize(update_deepmd_input(config, warning=False), multi_task=multitask)
+
+
+def _record_training(trainer: Trainer, *, data_rng_draws: int = 0) -> list:
+    """Record real optimizer inputs, excluding diagnostic reads of the data."""
+    records = []
+    get_data = trainer.get_data
+    train_step = trainer.train_step
+    in_train_step = False
+
+    def record_data(is_train: bool = True, task_key: str = "Default") -> tuple:
+        inputs, labels = get_data(is_train=is_train, task_key=task_key)
+        if is_train and in_train_step:
+            if data_rng_draws:
+                dp_random.random(data_rng_draws)
+            records.append(
+                (
+                    task_key,
+                    {
+                        key: value.detach().cpu().clone()
+                        for key, value in inputs.items()
+                        if torch.is_tensor(value)
+                    },
+                    labels["energy"].detach().cpu().clone(),
+                )
+            )
+        return inputs, labels
+
+    def record_step(task, step):
+        nonlocal in_train_step
+        in_train_step = True
+        try:
+            return train_step(task, step)
+        finally:
+            in_train_step = False
+
+    trainer.get_data = record_data
+    trainer.train_step = record_step
+    trainer.run()
+    return records
+
+
+def _assert_records_equal(expected: list, actual: list) -> None:
+    assert len(expected) == len(actual)
+    for (task, inputs, energy), (other_task, other_inputs, other_energy) in zip(
+        expected, actual, strict=True
+    ):
+        assert task == other_task
+        assert inputs.keys() == other_inputs.keys()
+        for key, value in inputs.items():
+            torch.testing.assert_close(value, other_inputs[key], rtol=0, atol=0)
+        torch.testing.assert_close(energy, other_energy, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("mode", ["directory", "hdf5", "grouped"])
+def test_setup_statistics_preserve_training_progress(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mode: str
+) -> None:
+    """Fresh setup statistics consume samples without consuming training state."""
+    monkeypatch.chdir(tmp_path)
+    path = tmp_path / ("data.h5" if mode == "hdf5" else "data")
+    if mode == "hdf5":
+        _write_hdf5(path)
+    else:
+        _write_directory(path)
+    system = str(path) + ("#/system" if mode == "hdf5" else "")
+    config = _config(system, mode="directory")
+    if mode == "grouped":
+        second = tmp_path / "second"
+        _write_directory(second)
+        for task in config["training"]["data_dict"].values():
+            task["training_data"] = [
+                {"systems": [system], "batch_size": 1, "prob": 0.4},
+                {"systems": [str(second)], "batch_size": 1, "prob": 0.6},
+            ]
+    for model in config["model"]["model_dict"].values():
+        model["data_stat_nbatch"] = 3
+    snapshots = []
+    make_stat_input = training_module.make_stat_input
+
+    def record_statistics(data: Any, nbatches: int, filter_enabled: bool) -> list[dict]:
+        before = collect_training_data_progress({"statistics": data})
+        assert "statistics" in before["tasks"]
+        sampled = make_stat_input(data, nbatches, filter_enabled)
+        assert collect_training_data_progress({"statistics": data}) != before
+        snapshots.append((data, before))
+        return sampled
+
+    monkeypatch.setattr(training_module, "make_stat_input", record_statistics)
+    trainer = get_trainer(config)
+    try:
+        assert len(snapshots) == len(trainer.training_data_by_task)
+        for data, expected in snapshots:
+            assert collect_training_data_progress({"statistics": data}) == expected
+    finally:
+        trainer._close_data_systems()
+
+
+@pytest.mark.parametrize(
+    ("mode", "seed"),
+    [
+        ("directory", 19),
+        ("lmdb", 19),
+        ("ragged", 19),
+        ("grouped", 19),
+        ("grouped", None),
+    ],
+)
+def test_restart_matches_uninterrupted_training(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mode: str, seed: int | None
+) -> None:
+    """Real training matches across passes, display, validation and setup."""
+    monkeypatch.chdir(tmp_path)
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    path = tmp_path / ("data" if mode in {"directory", "grouped"} else "data.lmdb")
+    if mode in {"directory", "grouped"}:
+        _write_directory(path)
+    else:
+        _write_lmdb(path, mixed=mode == "ragged")
+    config = _config(path, mode="directory" if mode == "grouped" else mode)
+    config["training"]["seed"] = seed
+    if mode == "grouped":
+        lmdb_paths = [tmp_path / f"part_{index}.lmdb" for index in range(2)]
+        for lmdb_path in lmdb_paths:
+            _write_lmdb(lmdb_path, mixed=False)
+        for task in config["training"]["data_dict"].values():
+            task["training_data"] = [
+                {"systems": [str(path)], "batch_size": 1, "prob": 0.4},
+                {
+                    "systems": [str(item) for item in lmdb_paths],
+                    "batch_size": "mix:7",
+                    "auto_prob": "prob_uniform",
+                    "prob": 0.6,
+                },
+            ]
+    dp_random.seed(19)
+    uninterrupted = get_trainer(deepcopy(config))
+    if mode == "ragged":
+        assert all(
+            source._reader.ragged_batches
+            for source in uninterrupted.training_data_by_task.values()
+        )
+    expected = _record_training(uninterrupted)
+    checkpoint = tmp_path / "model.ckpt-10.pt"
+    state = torch.load(checkpoint, weights_only=True)
+    assert DATA_PROGRESS_CHECKPOINT_KEY in state
+    # Restart construction and fresh validation readers may consume different
+    # random draws; none may leak into the resumed training stream.
+    dp_random.seed(9876)
+    restarted = get_trainer(deepcopy(config), restart_model=str(checkpoint))
+    assert restarted.start_step == 10
+    actual = _record_training(restarted)
+    _assert_records_equal(expected[10:], actual)
+    for key, value in uninterrupted._unwrapped.state_dict().items():
+        if torch.is_tensor(value):
+            torch.testing.assert_close(
+                value, restarted._unwrapped.state_dict()[key], rtol=1e-12, atol=1e-12
+            )
+
+
+def test_best_validation_checkpoint_resumes_training(tmp_path, monkeypatch) -> None:
+    """The real full-validation callback writes a restartable snapshot."""
+    monkeypatch.chdir(tmp_path)
+    path = tmp_path / "data"
+    _write_directory(path)
+    config = _config(path, mode="directory", multitask=False)
+    config["training"].update(numb_steps=6, save_freq=0)
+    config["validating"] = {
+        "full_validation": True,
+        "full_val_start": 3,
+        "validation_freq": 3,
+        "save_best": True,
+        "save_best_dir": str(tmp_path / "best"),
+        "max_best_ckpt": 2,
+    }
+    original = get_trainer(deepcopy(config))
+    expected = _record_training(original)
+    checkpoint = next((tmp_path / "best").glob("*-3.t-*.pt"))
+    restarted = get_trainer(deepcopy(config), restart_model=str(checkpoint))
+    assert restarted.start_step == 3
+    _assert_records_equal(expected[3:], _record_training(restarted))
+
+
+@pytest.mark.parametrize("reset", [False, True])
+def test_legacy_init_finetune_and_deployment(
+    tmp_path, monkeypatch, caplog, reset
+) -> None:
+    """Only a restart restores progress; deployment snapshots omit it."""
+    monkeypatch.chdir(tmp_path)
+    path = tmp_path / "data"
+    _write_directory(path)
+    config = _config(path, mode="directory", multitask=False)
+    config["training"].update(
+        numb_steps=3, save_freq=3, enable_ema=True, reset_data_progress=reset
+    )
+    trainer = get_trainer(deepcopy(config))
+    assert trainer.start_step == 0
+    trainer.run()
+    checkpoint = tmp_path / "model.ckpt-3.pt"
+    state = torch.load(checkpoint, weights_only=True)
+    assert DATA_PROGRESS_CHECKPOINT_KEY in state
+    ema_state = torch.load(tmp_path / "model_ema.ckpt-3.pt", weights_only=True)
+    assert DATA_PROGRESS_CHECKPOINT_KEY not in ema_state
+    assert "optimizer" not in ema_state
+    legacy = tmp_path / "legacy.pt"
+    state.pop(DATA_PROGRESS_CHECKPOINT_KEY)
+    torch.save(state, legacy)
+    restarted = get_trainer(deepcopy(config), restart_model=str(legacy))
+    assert restarted.start_step == 3
+    restarted._close_data_systems()
+    # A deliberately unsupported version proves these paths do not restore it.
+    state[DATA_PROGRESS_CHECKPOINT_KEY] = {"version": -1}
+    torch.save(state, checkpoint)
+    for option in ("init_model", "finetune_model"):
+        fresh = get_trainer(deepcopy(config), **{option: str(checkpoint)})
+        assert fresh.start_step == 0
+        fresh._close_data_systems()
+    assert "training.reset_data_progress is true" not in caplog.text
+
+
+@pytest.mark.parametrize("mode", ["directory", "hdf5", "lmdb"])
+def test_relocated_restart_can_reset_only_data_progress(
+    tmp_path, monkeypatch, caplog, mode
+) -> None:
+    """Opting out after a move has exactly the legacy restart semantics."""
+    monkeypatch.chdir(tmp_path)
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    if mode == "lmdb":
+        path = tmp_path / "data.lmdb"
+        _write_lmdb(path, mixed=False)
+    elif mode == "hdf5":
+        path = tmp_path / "data.h5"
+        _write_hdf5(path)
+    else:
+        path = tmp_path / "data"
+        _write_directory(path)
+    system = str(path) + ("#/system" if mode == "hdf5" else "")
+    config_mode = "lmdb" if mode == "lmdb" else "directory"
+    config = _config(system, mode=config_mode)
+    config["training"].update(numb_steps=3, save_freq=3, enable_ema=True)
+    original = get_trainer(deepcopy(config))
+    original.run()
+    checkpoint = tmp_path / "model.ckpt-3.pt"
+    state = torch.load(checkpoint, weights_only=True)
+    # Native restart maps every optimizer tensor, including Adam's step.
+    checkpoint_optimizer = torch.load(
+        checkpoint, map_location=DEVICE, weights_only=True
+    )["optimizer"]["state"]
+    assert DATA_PROGRESS_CHECKPOINT_KEY in state
+    legacy_state = deepcopy(state)
+    del legacy_state[DATA_PROGRESS_CHECKPOINT_KEY]
+    legacy_path = tmp_path / "legacy.pt"
+    torch.save(legacy_state, legacy_path)
+
+    relocated = path.with_name("relocated" + path.suffix)
+    if mode == "hdf5":
+        copyfile(path, relocated)
+    else:
+        copytree(path, relocated)
+    system = str(relocated) + ("#/system" if mode == "hdf5" else "")
+    config = _config(system, mode=config_mode)
+    config["training"].update(numb_steps=9, save_freq=3, enable_ema=True)
+    assert config["training"]["reset_data_progress"] is False
+    with pytest.raises(ValueError, match=r"[Pp]rogress"):
+        get_trainer(deepcopy(config), restart_model=str(checkpoint))
+    assert "training.reset_data_progress is true" not in caplog.text
+
+    dp_random.seed(9876)
+    legacy = get_trainer(deepcopy(config), restart_model=str(legacy_path))
+    expected_lr = legacy.scheduler.get_last_lr()
+    expected_optimizer = deepcopy(legacy.optimizer.state_dict())
+    expected_progress = collect_training_data_progress(legacy.training_data_by_task)
+    expected = _record_training(legacy)
+    dp_random.seed(9876)
+    config["training"]["reset_data_progress"] = True
+    reset = get_trainer(deepcopy(config), restart_model=str(checkpoint))
+    assert "training.reset_data_progress is true" in caplog.text
+    assert "Batches may repeat" in caplog.text
+    assert "exact data-order continuation is unavailable" in caplog.text
+    assert reset.start_step == legacy.start_step == 3
+    assert reset.scheduler.get_last_lr() == expected_lr
+    assert (
+        collect_training_data_progress(reset.training_data_by_task) == expected_progress
+    )
+    torch.testing.assert_close(reset.optimizer.state_dict(), expected_optimizer)
+    torch.testing.assert_close(
+        reset.optimizer.state_dict()["state"], checkpoint_optimizer
+    )
+    torch.testing.assert_close(reset.model_ema.state_dict(), state[EMA_CHECKPOINT_KEY])
+    for key, value in state["model"].items():
+        if torch.is_tensor(value):
+            torch.testing.assert_close(value, reset._unwrapped.state_dict()[key])
+    _assert_records_equal(expected, _record_training(reset))
+    assert reset.scheduler.get_last_lr() == legacy.scheduler.get_last_lr()
+    torch.testing.assert_close(
+        reset.optimizer.state_dict(), legacy.optimizer.state_dict()
+    )
+    torch.testing.assert_close(
+        reset.model_ema.state_dict(), legacy.model_ema.state_dict()
+    )
+    for key, value in legacy._unwrapped.state_dict().items():
+        if torch.is_tensor(value):
+            torch.testing.assert_close(
+                value, reset._unwrapped.state_dict()[key], rtol=1e-12, atol=1e-12
+            )
+
+
+@pytest.mark.parametrize("mode", ["directory", "hdf5", "lmdb"])
+def test_configured_seed_is_format_independent(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mode: str
+) -> None:
+    """Fresh streams are seeded consistently; exact restarts reject seed changes."""
+    monkeypatch.chdir(tmp_path)
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    path = (
+        tmp_path / {"directory": "data", "hdf5": "data.h5", "lmdb": "data.lmdb"}[mode]
+    )
+    if mode == "lmdb":
+        _write_lmdb(path, mixed=False)
+    elif mode == "hdf5":
+        _write_hdf5(path)
+    else:
+        _write_directory(path)
+    system = str(path) + ("#/system" if mode == "hdf5" else "")
+    config = _config(
+        system, mode="lmdb" if mode == "lmdb" else "directory", multitask=False
+    )
+    config["training"].update(seed=7, numb_steps=1, save_freq=1)
+    states = []
+    for external_seed in (1, 999):
+        dp_random.seed(external_seed)
+        trainer = get_trainer(deepcopy(config))
+        try:
+            states.append(collect_training_data_progress(trainer.training_data_by_task))
+        finally:
+            trainer._close_data_systems()
+    assert states[0] == states[1]
+
+    original = get_trainer(deepcopy(config))
+    original.run()
+    checkpoint = str(tmp_path / "model.ckpt-1.pt")
+    config["training"]["seed"] = 12
+    with pytest.raises(ValueError, match=r"training.seed changed"):
+        get_trainer(deepcopy(config), restart_model=checkpoint)
+    config["training"]["reset_data_progress"] = True
+    restarted = get_trainer(config, restart_model=checkpoint)
+    try:
+        assert restarted.start_step == 1
+        assert restarted._collect_data_progress()["configured_seed"] == 12
+    finally:
+        restarted._close_data_systems()
+
+
+@pytest.mark.parametrize("mode", ["directory", "lmdb"])
+def test_changed_world_size_retains_legacy_optimizer_restart(
+    tmp_path, monkeypatch, caplog, mode
+) -> None:
+    """A changed topology keeps native restart state, but not saved data order."""
+    monkeypatch.chdir(tmp_path)
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    path = tmp_path / ("data" if mode == "directory" else "data.lmdb")
+    if mode == "directory":
+        _write_directory(path)
+    else:
+        _write_lmdb(path, mixed=False)
+    config = _config(path, mode=mode, multitask=False)
+    config["training"].update(numb_steps=3, save_freq=3)
+    original = get_trainer(deepcopy(config))
+    original.run()
+    checkpoint = tmp_path / "model.ckpt-3.pt"
+    state = torch.load(checkpoint, weights_only=True)
+    checkpoint_optimizer = torch.load(
+        checkpoint, map_location=DEVICE, weights_only=True
+    )["optimizer"]["state"]
+    progress = state.pop(DATA_PROGRESS_CHECKPOINT_KEY)
+    legacy_path = tmp_path / "legacy.pt"
+    torch.save(state, legacy_path)
+    # Exercise the topology fallback at the native checkpoint-load boundary.
+    # The real 1-to-2-rank integration case remains in TestDDPRestart.
+    progress.update(world_size=2, ranks=progress["ranks"] * 2)
+    state[DATA_PROGRESS_CHECKPOINT_KEY] = progress
+    changed_path = tmp_path / "changed-world.pt"
+    torch.save(state, changed_path)
+    config["training"]["numb_steps"] = 6
+
+    dp_random.seed(9876)
+    legacy = get_trainer(deepcopy(config), restart_model=str(legacy_path))
+    expected_lr = legacy.scheduler.get_last_lr()
+    expected = _record_training(legacy)
+    dp_random.seed(9876)
+    changed = get_trainer(deepcopy(config), restart_model=str(changed_path))
+    assert "world size changed from 2 to 1" in caplog.text
+    assert "exact data-order continuation is unavailable" in caplog.text
+    assert changed.start_step == 3
+    assert changed.scheduler.get_last_lr() == expected_lr
+    torch.testing.assert_close(
+        changed.optimizer.state_dict()["state"], checkpoint_optimizer
+    )
+    for key, value in state["model"].items():
+        if torch.is_tensor(value):
+            torch.testing.assert_close(value, changed._unwrapped.state_dict()[key])
+    _assert_records_equal(expected, _record_training(changed))
+    for key, value in legacy._unwrapped.state_dict().items():
+        if torch.is_tensor(value):
+            torch.testing.assert_close(value, changed._unwrapped.state_dict()[key])
+
+
+def test_chief_display_preserves_training_progress(tmp_path, monkeypatch) -> None:
+    """Chief-only inactive-task and validation reads cannot move rank cursors."""
+    monkeypatch.chdir(tmp_path)
+    path = tmp_path / "data"
+    _write_directory(path)
+    trainer = get_trainer(_config(path, mode="directory"))
+    active = trainer.training_tasks["first"]
+    result = trainer.train_step(active, 0)
+    before = collect_training_data_progress(trainer.training_data_by_task)
+    trainer.collect_display_results(
+        trainer.training_tasks, active_task=active, step=0, step_result=result
+    )
+    assert collect_training_data_progress(trainer.training_data_by_task) == before
+    trainer._close_data_systems()
+
+
+def test_ddp_warmup_preserves_training_progress(tmp_path, monkeypatch) -> None:
+    """Exercise the warmup body on CPU without invoking a compiler or NCCL."""
+    monkeypatch.chdir(tmp_path)
+    path = tmp_path / "data"
+    _write_directory(path)
+    trainer = get_trainer(_config(path, mode="directory"))
+    before = collect_training_data_progress(trainer.training_data_by_task)
+
+    class DDPWrapper:
+        def __init__(self, module):
+            self.module = module
+
+    trainer.wrapper = DDPWrapper(trainer.wrapper)
+    trainer.enable_compile = True
+    with (
+        patch.object(torch.nn.parallel, "DistributedDataParallel", DDPWrapper),
+        patch.object(dist, "is_initialized", return_value=True),
+        patch.object(dist, "get_world_size", return_value=2),
+        patch.object(dist.distributed_c10d, "_get_default_store") as store,
+    ):
+        store.return_value.add.return_value = 2
+        trainer._precompile_outside_collectives()
+    assert collect_training_data_progress(trainer.training_data_by_task) == before
+    trainer._close_data_systems()
+
+
+def test_full_validation_checkpoint_uses_training_rng() -> None:
+    """A best-checkpoint callback runs before validation returns to the loop."""
+    trainer = Trainer.__new__(Trainer)
+    dp_random.seed(42)
+    before = dp_random.get_state()
+    saved = []
+
+    def save(*args, **kwargs):
+        saved.append(dp_random.get_state())
+
+    def validate(*, save_checkpoint, **kwargs):
+        dp_random.random(100)
+        save_checkpoint(Path("best.pt"), lr=1e-3, step=1)
+        dp_random.random(100)
+
+    trainer._collect_data_progress = Mock(return_value={"saved_rng": before})
+    trainer.full_validator = Mock(run=validate, save_best=True)
+    trainer.ema_full_validator = None
+    trainer._save_full_validation_checkpoint = save
+    trainer.run_full_validation(step=0, display_step=1, learning_rate=1e-3)
+    assert saved == [before]
+    assert dp_random.get_state() == before
+
+
+def _set_mock_rank(trainer: Trainer, rank: int) -> None:
+    trainer.is_distributed = True
+    trainer.rank = rank
+    trainer.world_size = 2
+    trainer.rank_context = RankContext(rank=rank, world_size=2)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_every_rank_joins_periodic_final_and_bias_saves(
+    tmp_path, monkeypatch, rank
+) -> None:
+    """Run the native loop; only checkpoint boundaries gather rank cursors."""
+    monkeypatch.chdir(tmp_path)
+    path = tmp_path / "data"
+    _write_directory(path)
+    config = _config(path, mode="directory", multitask=False)
+    config["training"].update(
+        numb_steps=3,
+        save_freq=2,
+        disp_training=False,
+        change_bias_after_training=True,
+        enable_ema=True,
+    )
+    trainer = get_trainer(config)
+    _set_mock_rank(trainer, rank)
+    trainer._change_bias_after_training = Mock()
+    gathered = []
+    collectives = []
+
+    def gather(ranks, local):
+        if isinstance(local, bool):
+            collectives.append("preflight")
+        else:
+            collectives.append("progress")
+            gathered.append(deepcopy(local))
+        ranks[:] = [deepcopy(local), deepcopy(local)]
+
+    monkeypatch.setattr(dist, "all_gather_object", gather)
+    monkeypatch.setattr(dist, "get_world_size", lambda: 2)
+    steps = []
+    save = trainer.save_checkpoint
+
+    def record_save(step):
+        steps.append(step)
+        save(step)
+
+    trainer.save_checkpoint = record_save
+    trainer.run()
+    assert steps == [2, 3, 3]
+    assert len(gathered) == 3  # No per-batch or EMA-snapshot gather.
+    assert collectives == ["preflight", "progress"] * 3
+    trainer._change_bias_after_training.assert_called_once()
+    checkpoint = tmp_path / "model.ckpt-3.pt"
+    assert checkpoint.exists() == (rank == 0)
+    if rank == 0:
+        progress = torch.load(checkpoint, weights_only=True)[
+            DATA_PROGRESS_CHECKPOINT_KEY
+        ]
+        assert progress == {
+            "version": 1,
+            "world_size": 2,
+            "configured_seed": config["training"]["seed"],
+            "task_rng": trainer._task_rng.bit_generator.state,
+            "ranks": [gathered[-1]] * 2,
+        }
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+@pytest.mark.parametrize("due,save_best", [(True, True), (True, False), (False, True)])
+def test_full_validation_precollects_before_chief_only_save(
+    tmp_path, monkeypatch, rank, due, save_best
+) -> None:
+    """The chief callback writes a complete envelope without a nested gather."""
+    monkeypatch.chdir(tmp_path)
+    path = tmp_path / "data"
+    _write_directory(path)
+    config = _config(path, mode="directory", multitask=False)
+    config["training"]["enable_ema"] = True
+    trainer = get_trainer(config)
+    _set_mock_rank(trainer, rank)
+    before = collect_training_data_progress(trainer.training_data_by_task)
+    gathered = []
+    calls = []
+
+    def gather(ranks, local):
+        assert not calls
+        gathered.append(deepcopy(local))
+        ranks[:] = [deepcopy(local), deepcopy(local)]
+
+    monkeypatch.setattr(dist, "all_gather_object", gather)
+
+    def validate(*, save_checkpoint, **kwargs):
+        calls.append("live")
+        assert len(gathered) == int(save_best)
+        dp_random.random(30)
+        if rank == 0 and save_best:
+            save_checkpoint(tmp_path / "best.pt", step=1)
+
+    def validate_ema(*, save_checkpoint, **kwargs):
+        calls.append("ema")
+        if rank == 0 and save_best:
+            save_checkpoint(tmp_path / "best_ema.pt", step=1)
+
+    trainer.full_validator = Mock(
+        should_run=Mock(return_value=due), save_best=save_best, run=validate
+    )
+    trainer.ema_full_validator = Mock(
+        should_run=Mock(return_value=due), save_best=save_best, run=validate_ema
+    )
+    trainer.run_full_validation(step=0, display_step=1, learning_rate=1e-3)
+    assert len(gathered) == int(due and save_best)
+    assert calls == (["live", "ema"] if due else [])
+    assert dp_random.get_state() == before["rng"]
+    checkpoint = tmp_path / "best.pt"
+    assert checkpoint.exists() == (rank == 0 and due and save_best)
+    if checkpoint.exists():
+        state = torch.load(checkpoint, weights_only=True)
+        assert state[DATA_PROGRESS_CHECKPOINT_KEY] == {
+            "version": 1,
+            "world_size": 2,
+            "configured_seed": config["training"]["seed"],
+            "task_rng": trainer._task_rng.bit_generator.state,
+            "ranks": [before, before],
+        }
+        ema = torch.load(tmp_path / "best_ema.pt", weights_only=True)
+        assert DATA_PROGRESS_CHECKPOINT_KEY not in ema
+    trainer._close_data_systems()
+
+
+def test_filtered_rank_progress_restores_its_own_cursor(tmp_path, monkeypatch) -> None:
+    """Local filter retries make rank zero's cursor wrong for another rank."""
+    path = _create_lmdb(str(tmp_path / "filtered.lmdb"), nframes=12, natoms=2)
+
+    def make_rank(rank):
+        data = LmdbDataSystem(
+            path,
+            ["O", "H"],
+            batch_size=1,
+            seed=7,
+            num_workers=0,
+            rank=rank,
+            world_size=2,
+        )
+        data.add_data_requirements(
+            [
+                DataRequirementItem(
+                    "pair_margin", 1, default=2.0, source_policy="derived"
+                )
+            ]
+        )
+        trainer = Trainer.__new__(Trainer)
+        _set_mock_rank(trainer, rank)
+        trainer._configured_data_seed = 7
+        trainer._task_rng = np.random.default_rng(7)
+        trainer.training_data_by_task = {"Default": data}
+        trainer.pair_filter_enabled_by_task = {"Default": True}
+        return trainer, data
+
+    originals = [make_rank(rank) for rank in range(2)]
+    for trainer, data in originals:
+        trainer._next_training_batch(data, "Default")
+    progress = [
+        collect_training_data_progress(trainer.training_data_by_task)
+        for trainer, _ in originals
+    ]
+    assert progress[0]["tasks"] != progress[1]["tasks"]
+
+    def gather(ranks, local):
+        ranks[:] = deepcopy(progress)
+
+    monkeypatch.setattr(dist, "all_gather_object", gather)
+    checkpoint = tmp_path / "progress.pt"
+    torch.save(originals[0][0]._collect_data_progress(), checkpoint)
+    envelope = torch.load(checkpoint, weights_only=True)
+    for rank, (trainer, data) in enumerate(originals):
+        expected = [
+            trainer._next_training_batch(data, "Default")["coord"] for _ in range(4)
+        ]
+        restored, restored_data = make_rank(rank)
+        restored._restore_data_progress(envelope)
+        actual = [
+            restored._next_training_batch(restored_data, "Default")["coord"]
+            for _ in range(4)
+        ]
+        for want, got in zip(expected, actual, strict=True):
+            np.testing.assert_array_equal(want, got)
+        data.close()
+        restored_data.close()
+
+
+@pytest.mark.parametrize(
+    "saved_size,current_size,rank", [(1, 2, 0), (1, 2, 1), (2, 1, 0), (2, 3, 2)]
+)
+def test_changed_world_size_leaves_data_and_rng_untouched(
+    caplog, saved_size, current_size, rank
+) -> None:
+    """Topology fallback neither remaps a cursor nor restores another rank's RNG."""
+    trainer = Trainer.__new__(Trainer)
+    trainer.rank = rank
+    trainer.world_size = current_size
+    trainer._task_rng = np.random.default_rng(7)
+    task_rng_state = deepcopy(trainer._task_rng.bit_generator.state)
+    source = Mock()
+    source.state_dict.return_value = {"cursor": 11}
+    trainer.training_data_by_task = {"Default": source}
+    local = collect_training_data_progress(trainer.training_data_by_task)
+    local["tasks"]["Default"] = {"cursor": 99}
+    progress = {"version": 1, "world_size": saved_size, "ranks": [local] * saved_size}
+    dp_random.random(5)
+    before = collect_training_data_progress(trainer.training_data_by_task)
+    with patch(
+        "deepmd.pt_expt.train.training.restore_training_data_progress"
+    ) as restore:
+        trainer._restore_data_progress(progress)
+    restore.assert_not_called()
+    source.load_state_dict.assert_not_called()
+    assert collect_training_data_progress(trainer.training_data_by_task) == before
+    assert trainer._task_rng.bit_generator.state == task_rng_state
+    assert ("Checkpoint world size changed" in caplog.text) == (rank == 0)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_explicit_reset_leaves_every_task_and_rng_untouched(caplog, rank) -> None:
+    """An explicit reset does not even invoke a task or shared RNG restore."""
+    trainer = Trainer.__new__(Trainer)
+    _set_mock_rank(trainer, rank)
+    trainer._task_rng = np.random.default_rng(7)
+    task_rng_state = deepcopy(trainer._task_rng.bit_generator.state)
+    sources = {"first": Mock(), "second": Mock()}
+    for index, source in enumerate(sources.values()):
+        source.state_dict.return_value = {"cursor": index + 10}
+    trainer.training_data_by_task = sources
+    dp_random.seed(42)
+    local = collect_training_data_progress(sources)
+    for task in local["tasks"]:
+        local["tasks"][task] = {"cursor": 99}
+    progress = {"version": 1, "world_size": 2, "ranks": [local, local]}
+    dp_random.random(5)
+    before = collect_training_data_progress(sources)
+    with (
+        patch(
+            "deepmd.pt_expt.train.training.restore_training_data_progress"
+        ) as restore,
+        patch.object(dp_random, "set_state") as set_rng,
+    ):
+        trainer._restore_data_progress(progress, reset=True)
+    restore.assert_not_called()
+    set_rng.assert_not_called()
+    for source in sources.values():
+        source.load_state_dict.assert_not_called()
+    assert collect_training_data_progress(sources) == before
+    assert trainer._task_rng.bit_generator.state == task_rng_state
+    assert ("training.reset_data_progress is true" in caplog.text) == (rank == 0)
+    assert "Checkpoint world size changed" not in caplog.text
+
+
+@pytest.mark.parametrize("reset", [False, True])
+@pytest.mark.parametrize("world_size", [1, 2])
+@pytest.mark.parametrize(
+    "changed",
+    [
+        {"version": 2},
+        {"world_size": 3},
+        {"world_size": 0},
+        {"world_size": -1},
+        {"world_size": True},
+        {"world_size": 2.0},
+        {"world_size": "2"},
+        {"ranks": []},
+        {"ranks": {}},
+        {"ranks": [None, None]},
+    ],
+)
+def test_malformed_rank_envelope_is_rejected(
+    changed, world_size, caplog, reset
+) -> None:
+    trainer = Trainer.__new__(Trainer)
+    trainer.rank = 0
+    trainer.world_size = world_size
+    trainer.training_data_by_task = {}
+    local = collect_training_data_progress({})
+    progress = {"version": 1, "world_size": 2, "ranks": [local, local]}
+    with pytest.raises(ValueError, match=r"[Dd]ata-progress"):
+        trainer._restore_data_progress(progress | changed, reset=reset)
+    assert "Checkpoint world size changed" not in caplog.text
+    assert "training.reset_data_progress is true" not in caplog.text
+    trainer._restore_data_progress(None, reset=reset)
+    assert "training.reset_data_progress is true" not in caplog.text
+
+
+def _ddp_restart_worker(rank: int, directory: str, rendezvous: str, mode: str) -> None:
+    torch.set_num_threads(1)
+    root = Path(directory)
+    dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=2)
+    try:
+        path = root / ("data" if mode == "directory" else "data.lmdb")
+        config = _config(path, mode=mode)
+        config["training"].update(
+            save_ckpt=str(root / "model.ckpt"),
+            disp_file=str(root / "lcurve.out"),
+        )
+        if mode == "directory":
+            unseeded_config = deepcopy(config)
+            unseeded_config["training"]["seed"] = None
+            unseeded = get_trainer(unseeded_config)
+            choices = [
+                unseeded.select_task(unseeded.training_tasks).key for _ in range(8)
+            ]
+            ranks = [None, None]
+            dist.all_gather_object(ranks, choices)
+            assert ranks[0] == ranks[1]
+            unseeded._close_data_systems()
+        dp_random.seed(19)
+        trainer = get_trainer(deepcopy(config))
+        with preserve_training_data_progress(trainer.training_data_by_task):
+            first = next(iter(trainer.training_data_by_task.values()))
+            frames = [np.asarray(first.get_batch()["coord"]).tolist() for _ in range(8)]
+        ranks = [None, None]
+        dist.all_gather_object(ranks, frames)
+        assert ranks[0] != ranks[1]
+        expected = _record_training(trainer, data_rng_draws=rank + 1)
+        dist.all_gather_object(ranks, [record[0] for record in expected])
+        assert ranks[0] == ranks[1]
+        dist.barrier()
+        dp_random.seed(9876)
+        restarted = get_trainer(
+            deepcopy(config), restart_model=str(root / "model.ckpt-10.pt")
+        )
+        actual = _record_training(restarted, data_rng_draws=rank + 1)
+        _assert_records_equal(expected[10:], actual)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(not dist.is_gloo_available(), reason="Requires Gloo")
+@pytest.mark.parametrize("mode", ["directory", "lmdb"])
+def test_two_rank_restart_preserves_each_rank(tmp_path, monkeypatch, mode) -> None:
+    """Chief-only statistics/display do not desynchronize saved rank cursors."""
+    monkeypatch.setenv("DEVICE", "cpu")
+    monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0")
+    if mode == "directory":
+        _write_directory(tmp_path / "data")
+    else:
+        _write_lmdb(tmp_path / "data.lmdb", mixed=False)
+    mp.spawn(
+        _ddp_restart_worker,
+        args=(str(tmp_path), (tmp_path / "rendezvous").as_uri(), mode),
+        nprocs=2,
+        join=True,
+    )
diff --git a/source/tests/pt_expt/test_dp_freeze.py b/source/tests/pt_expt/test_dp_freeze.py
index cc7a0bd265..51aabc953b 100644
--- a/source/tests/pt_expt/test_dp_freeze.py
+++ b/source/tests/pt_expt/test_dp_freeze.py
@@ -99,12 +99,6 @@ def setUpClass(cls) -> None:
     def tearDownClass(cls) -> None:
         shutil.rmtree(cls.tmpdir)
 
-    def test_freeze_pte(self) -> None:
-        """Freeze to .pte and verify the file is created."""
-        output = os.path.join(self.tmpdir, "frozen_model.pte")
-        freeze(model=self.ckpt_file, output=output)
-        self.assertTrue(os.path.exists(output))
-
     def test_freeze_main_dispatcher(self) -> None:
         """Test main() CLI dispatcher with freeze command."""
         output_file = os.path.join(self.tmpdir, "frozen_via_main.pte")
@@ -226,34 +220,6 @@ def test_freeze_graph_dpa1(self) -> None:
             metadata = json.loads(zf.read(meta_name))
         self.assertEqual(metadata["lower_input_kind"], "graph")
 
-    def test_freeze_pt2(self) -> None:
-        """Freeze to .pt2 (AOTInductor) and verify the file is loadable."""
-        self.assertTrue(os.path.exists(self.shared_pt2))
-
-        # Verify the .pt2 can be loaded and evaluated via DeepPot
-        import numpy as np
-
-        from deepmd.infer import (
-            DeepPot,
-        )
-
-        dp = DeepPot(self.shared_pt2)
-        self.assertEqual(dp.get_type_map(), ["O", "H", "B"])
-        rcut = dp.get_rcut()
-        self.assertGreater(rcut, 0.0)
-
-        # Quick smoke-test eval
-        coord = np.array(
-            [0.0, 0.0, 0.1, 1.0, 0.3, 0.2, 0.1, 1.9, 3.4],
-            dtype=np.float64,
-        )
-        box = np.array([5.0, 0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 5.0], dtype=np.float64)
-        atype = [0, 1, 2]
-        e, f, v = dp.eval(coord, box, atype)
-        self.assertEqual(e.shape, (1, 1))
-        self.assertEqual(f.shape, (1, 3, 3))
-        self.assertEqual(v.shape, (1, 9))
-
     def test_freeze_pt2_eval_consistency(self) -> None:
         """Verify .pte and .pt2 produce identical results."""
         import numpy as np
@@ -269,11 +235,18 @@ def test_freeze_pt2_eval_consistency(self) -> None:
         box = np.array([5.0, 0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 5.0], dtype=np.float64)
         atype = [0, 1, 2]
 
+        self.assertTrue(os.path.exists(self.shared_pte))
+        self.assertTrue(os.path.exists(self.shared_pt2))
         dp_pte = DeepPot(self.shared_pte)
         dp_pt2 = DeepPot(self.shared_pt2)
+        self.assertEqual(dp_pt2.get_type_map(), ["O", "H", "B"])
+        self.assertGreater(dp_pt2.get_rcut(), 0.0)
 
         e_pte, f_pte, v_pte = dp_pte.eval(coord, box, atype)
         e_pt2, f_pt2, v_pt2 = dp_pt2.eval(coord, box, atype)
+        self.assertEqual(e_pt2.shape, (1, 1))
+        self.assertEqual(f_pt2.shape, (1, 3, 3))
+        self.assertEqual(v_pt2.shape, (1, 9))
 
         np.testing.assert_allclose(e_pte, e_pt2, atol=1e-10)
         np.testing.assert_allclose(f_pte, f_pt2, atol=1e-10)
@@ -319,9 +292,6 @@ def test_nonspin_model_rejects_spin(self) -> None:
             DeepPot,
         )
 
-        pt2_path = os.path.join(self.tmpdir, "nonspin_reject.pt2")
-        freeze(model=self.ckpt_file, output=pt2_path)
-
         coord = np.array(
             [0.0, 0.0, 0.1, 1.0, 0.3, 0.2, 0.1, 1.9, 3.4],
             dtype=np.float64,
@@ -330,7 +300,7 @@ def test_nonspin_model_rejects_spin(self) -> None:
         atype = [0, 1, 2]
         spin = np.zeros(9, dtype=np.float64)
 
-        dp = DeepPot(pt2_path)
+        dp = DeepPot(self.shared_pt2)
         with self.assertRaises(ValueError):
             dp.eval(coord, box, atype, spin=spin)
 
diff --git a/source/tests/pt_expt/test_finetune.py b/source/tests/pt_expt/test_finetune.py
index 859e1df8c2..bd5e46a28f 100644
--- a/source/tests/pt_expt/test_finetune.py
+++ b/source/tests/pt_expt/test_finetune.py
@@ -7,7 +7,8 @@
 
 Part B: CLI end-to-end tests (TestFinetuneCLI)
   - test_finetune_cli
-  - test_finetune_cli_use_pretrain_script
+  - test_finetune_from_pte
+  - test_finetune_from_pt2
   - test_finetune_random_fitting
 """
 
@@ -599,124 +600,97 @@ def _assert_inherited_weights_match(
                             f"when random_fitting=True",
                         )
 
-    def test_finetune_cli(self) -> None:
-        """Train -> finetune via main() dispatcher -> verify checkpoint exists."""
+    def _check_finetune_artifact(self, suffix: str) -> None:
+        """Reuse one pretrained artifact for both CLI configuration policies."""
         from deepmd.pt_expt.entrypoints.main import (
-            main,
+            freeze,
         )
 
         tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_cli_")
         old_cwd = os.getcwd()
         os.chdir(tmpdir)
         try:
-            # Phase 1: train pretrained model
-            config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1)
-            config = update_deepmd_input(config, warning=False)
-            config = normalize(config)
-            ckpt_path = self._train_pretrained(config, tmpdir)
-
-            # Save original bias
-            state = torch.load(ckpt_path, map_location=DEVICE, weights_only=True)
-            model_state = state["model"] if "model" in state else state
-            original_model = get_model(model_state["_extra_state"]["model_params"]).to(
-                DEVICE
-            )
-            original_wrapper = ModelWrapper(original_model)
-            original_wrapper.load_state_dict(model_state)
-            original_bias = to_numpy_array(original_model.get_out_bias()).copy()
-
-            # Phase 2: finetune via CLI (lr=0 so weights stay unchanged)
-            ft_config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1)
-            ft_config["learning_rate"]["start_lr"] = 1e-30
-            ft_config["learning_rate"]["stop_lr"] = 1e-30
-            ft_config_file = os.path.join(tmpdir, "finetune_input.json")
-            with open(ft_config_file, "w") as f:
-                json.dump(ft_config, f)
-
-            main(
-                [
-                    "train",
-                    ft_config_file,
-                    "--finetune",
-                    ckpt_path,
-                    "--skip-neighbor-stat",
-                ]
-            )
-
-            # Verify new checkpoint exists
-            ft_ckpt = os.path.join(tmpdir, "model.ckpt.pt")
-            self.assertTrue(os.path.exists(ft_ckpt), "Finetune checkpoint not found")
-
-            # Load finetuned model and verify bias changed
-            ft_state = torch.load(ft_ckpt, map_location=DEVICE, weights_only=True)
-            ft_model_state = ft_state["model"] if "model" in ft_state else ft_state
-            ft_model = get_model(ft_model_state["_extra_state"]["model_params"]).to(
-                DEVICE
-            )
-            ft_wrapper = ModelWrapper(ft_model)
-            ft_wrapper.load_state_dict(ft_model_state)
-            ft_bias = to_numpy_array(ft_model.get_out_bias())
-
-            # Bias should have been adjusted (may or may not differ depending
-            # on data, but the checkpoint should at least be valid)
-            self.assertEqual(original_bias.shape, ft_bias.shape)
-
-            # Inherited weights (descriptor + fitting) must match pretrained.
-            # lr=0 so training step doesn't modify weights.
-            self._assert_inherited_weights_match(
-                ft_model_state, model_state, random_fitting=False
+            config = normalize(
+                update_deepmd_input(
+                    _make_config(self.data_dir, model_se_e2_a), warning=False
+                )
             )
+            checkpoint = self._train_pretrained(config, tmpdir)
+            pretrained = torch.load(checkpoint, map_location=DEVICE, weights_only=True)
+            pretrained = pretrained["model"] if "model" in pretrained else pretrained
+            artifact = os.path.join(tmpdir, "pretrained" + suffix)
+            if suffix == ".pt":
+                shutil.copy2(checkpoint, artifact)
+            else:
+                freeze(model=checkpoint, output=artifact)
+            self.assertTrue(os.path.exists(artifact))
+
+            for use_pretrain_script in (False, True):
+                with self.subTest(use_pretrain_script=use_pretrain_script):
+                    output_dir = os.path.join(tmpdir, str(use_pretrain_script))
+                    os.mkdir(output_dir)
+                    os.chdir(output_dir)
+                    self._run_finetune_cli(artifact, pretrained, use_pretrain_script)
         finally:
             os.chdir(old_cwd)
             shutil.rmtree(tmpdir, ignore_errors=True)
 
-    def test_finetune_cli_use_pretrain_script(self) -> None:
-        """Finetune with --use-pretrain-script -> config copied from pretrained."""
+    def _run_finetune_cli(
+        self, artifact: str, pretrained: dict, use_pretrain_script: bool
+    ) -> None:
+        """Fine-tune in an isolated output directory and verify the chosen policy."""
         from deepmd.pt_expt.entrypoints.main import (
             main,
         )
 
-        tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_pretrain_")
-        old_cwd = os.getcwd()
-        os.chdir(tmpdir)
-        try:
-            # Phase 1: train pretrained model
-            config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1)
-            config = update_deepmd_input(config, warning=False)
-            config = normalize(config)
-            ckpt_path = self._train_pretrained(config, tmpdir)
-
-            # Phase 2: finetune with --use-pretrain-script
-            # Use a config with different descriptor neuron sizes
-            ft_model_params = deepcopy(model_se_e2_a)
-            ft_model_params["descriptor"]["neuron"] = [4, 8]  # different
-            ft_config = _make_config(self.data_dir, ft_model_params, numb_steps=1)
-            ft_config_file = os.path.join(tmpdir, "finetune_input.json")
-            with open(ft_config_file, "w") as f:
-                json.dump(ft_config, f)
-
-            main(
-                [
-                    "train",
-                    ft_config_file,
-                    "--finetune",
-                    ckpt_path,
-                    "--use-pretrain-script",
-                    "--skip-neighbor-stat",
-                ]
-            )
-
-            # Verify the output config was updated from pretrained
-            with open(os.path.join(tmpdir, "out.json")) as f:
-                output_config = json.load(f)
-            # Descriptor neuron should be from pretrained, not from ft_config
+        model_params = deepcopy(model_se_e2_a)
+        if use_pretrain_script:
+            model_params["descriptor"]["neuron"] = [4, 8]
+        config = _make_config(self.data_dir, model_params)
+        if not use_pretrain_script:
+            config["learning_rate"]["start_lr"] = 1e-30
+            config["learning_rate"]["stop_lr"] = 1e-30
+        with open("finetune_input.json", "w") as handle:
+            json.dump(config, handle)
+        flags = ["train", "finetune_input.json", "--finetune", artifact]
+        if use_pretrain_script:
+            flags.append("--use-pretrain-script")
+        main([*flags, "--skip-neighbor-stat"])
+
+        self.assertTrue(os.path.exists("model.ckpt.pt"))
+        state = torch.load("model.ckpt.pt", map_location=DEVICE, weights_only=True)
+        state = state["model"] if "model" in state else state
+        self.assertIn("_extra_state", state)
+        if use_pretrain_script:
+            with open("out.json") as handle:
+                output_config = json.load(handle)
             self.assertEqual(
                 output_config["model"]["descriptor"]["neuron"],
                 model_se_e2_a["descriptor"]["neuron"],
             )
-        finally:
-            os.chdir(old_cwd)
-            shutil.rmtree(tmpdir, ignore_errors=True)
+        else:
+            self._assert_inherited_weights_match(state, pretrained)
+            biases = []
+            for model_state in (pretrained, state):
+                model = get_model(model_state["_extra_state"]["model_params"]).to(
+                    DEVICE
+                )
+                wrapper = ModelWrapper(model)
+                wrapper.load_state_dict(model_state)
+                biases.append(tuple(model.get_out_bias().shape))
+            self.assertEqual(*biases)
+
+    def test_finetune_cli(self) -> None:
+        """Fine-tune a checkpoint with explicit or pretrained model configuration."""
+        self._check_finetune_artifact(".pt")
+
+    def test_finetune_from_pte(self) -> None:
+        """Fine-tune a .pte artifact under both configuration policies."""
+        self._check_finetune_artifact(".pte")
+
+    def test_finetune_from_pt2(self) -> None:
+        """Fine-tune a .pt2 artifact under both configuration policies."""
+        self._check_finetune_artifact(".pt2")
 
     def test_finetune_random_fitting(self) -> None:
         """Finetune with --model-branch RANDOM -> descriptor from pretrained, fitting random."""
@@ -868,234 +842,75 @@ def test_finetune_random_fitting_into_bridged(self) -> None:
         """A RANDOM fitting of a bridged target takes the plain descriptor."""
         self._finetune_across_bridging(source_bridged=False, model_branch="RANDOM")
 
-    def test_finetune_from_pte(self) -> None:
-        """Train -> freeze to .pte -> finetune from .pte -> verify checkpoint."""
-        from deepmd.pt_expt.entrypoints.main import (
-            freeze,
-            main,
-        )
-
-        tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_pte_")
-        old_cwd = os.getcwd()
-        os.chdir(tmpdir)
-        try:
-            # Phase 1: train pretrained model
-            config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1)
-            config = update_deepmd_input(config, warning=False)
-            config = normalize(config)
-            ckpt_path = self._train_pretrained(config, tmpdir)
+    def test_finetune_short_range_into_long_range(self) -> None:
+        """A checkpoint without a long-range module fine-tunes into one with it.
 
-            # Phase 2: freeze to .pte
-            pte_path = os.path.join(tmpdir, "frozen.pte")
-            freeze(model=ckpt_path, output=pte_path)
-            self.assertTrue(os.path.exists(pte_path))
-
-            # Phase 3: finetune from .pte via CLI (lr=0 so weights stay unchanged)
-            ft_config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1)
-            ft_config["learning_rate"]["start_lr"] = 1e-30
-            ft_config["learning_rate"]["stop_lr"] = 1e-30
-            ft_config_file = os.path.join(tmpdir, "finetune_input.json")
-            with open(ft_config_file, "w") as f:
-                json.dump(ft_config, f)
-
-            main(
-                [
-                    "train",
-                    ft_config_file,
-                    "--finetune",
-                    pte_path,
-                    "--skip-neighbor-stat",
-                ]
-            )
-
-            # Verify new checkpoint exists
-            ft_ckpt = os.path.join(tmpdir, "model.ckpt.pt")
-            self.assertTrue(os.path.exists(ft_ckpt), "Finetune checkpoint not found")
-
-            # Load finetuned model and verify it's valid
-            ft_state = torch.load(ft_ckpt, map_location=DEVICE, weights_only=True)
-            ft_model_state = ft_state["model"] if "model" in ft_state else ft_state
-            self.assertIn("_extra_state", ft_model_state)
-
-            # Load pretrained from .pt for weight comparison
-            pre_state = torch.load(ckpt_path, map_location=DEVICE, weights_only=True)
-            pre_model_state = pre_state["model"] if "model" in pre_state else pre_state
-
-            # Inherited weights must match pretrained
-            self._assert_inherited_weights_match(
-                ft_model_state, pre_model_state, random_fitting=False
-            )
-        finally:
-            os.chdir(old_cwd)
-            shutil.rmtree(tmpdir, ignore_errors=True)
-
-    def test_finetune_from_pte_use_pretrain_script(self) -> None:
-        """Train -> freeze to .pte -> finetune with --use-pretrain-script."""
-        from deepmd.pt_expt.entrypoints.main import (
-            freeze,
-            main,
-        )
-
-        tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_pte_ups_")
+        The descriptor and the fitting net come from the checkpoint, while the
+        new periodic module starts from its initialization and takes the
+        input statistics of the fine-tuning data, here the bounds of its fixed
+        k grid, before the bias adjustment evaluates the complete model.
+        """
+        plain = {
+            "type_map": ["O", "H"],
+            "descriptor": {
+                "type": "dpa4c",
+                "rcut": 4.0,
+                "channels": 8,
+                "lmax": 2,
+                "n_radial": 4,
+                "seed": 1,
+            },
+            "fitting_net": {"neuron": [16, 16], "seed": 1},
+            "data_stat_nbatch": 1,
+        }
+        target = deepcopy(plain)
+        target["long_range"] = {
+            "type": "latent_charge",
+            "charge_net": {"neuron": [8], "seed": 1},
+            "kernel": {"type": "les", "alpha": 1.0},
+            "charge_constraint": {"source": "neutral"},
+            "periodic": True,
+            "k_grid_mode": "fixed",
+        }
+        tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_long_range_")
         old_cwd = os.getcwd()
         os.chdir(tmpdir)
         try:
-            # Phase 1: train pretrained model
-            config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1)
-            config = update_deepmd_input(config, warning=False)
-            config = normalize(config)
-            ckpt_path = self._train_pretrained(config, tmpdir)
-
-            # Phase 2: freeze to .pte (embeds model_params)
-            pte_path = os.path.join(tmpdir, "frozen.pte")
-            freeze(model=ckpt_path, output=pte_path)
-
-            # Phase 3: finetune from .pte with --use-pretrain-script
-            ft_model_params = deepcopy(model_se_e2_a)
-            ft_model_params["descriptor"]["neuron"] = [4, 8]  # different
-            ft_config = _make_config(self.data_dir, ft_model_params, numb_steps=1)
-            ft_config_file = os.path.join(tmpdir, "finetune_input.json")
-            with open(ft_config_file, "w") as f:
-                json.dump(ft_config, f)
-
-            main(
-                [
-                    "train",
-                    ft_config_file,
-                    "--finetune",
-                    pte_path,
-                    "--use-pretrain-script",
-                    "--skip-neighbor-stat",
-                ]
-            )
-
-            # Verify the output config was updated from pretrained
-            with open(os.path.join(tmpdir, "out.json")) as f:
-                output_config = json.load(f)
-            # Descriptor neuron should be from pretrained, not from ft_config
-            self.assertEqual(
-                output_config["model"]["descriptor"]["neuron"],
-                model_se_e2_a["descriptor"]["neuron"],
+            config = normalize(
+                update_deepmd_input(
+                    _make_config(self.data_dir, plain, numb_steps=1), warning=False
+                )
             )
-        finally:
-            os.chdir(old_cwd)
-            shutil.rmtree(tmpdir, ignore_errors=True)
-
-    def test_finetune_from_pt2(self) -> None:
-        """Train -> freeze to .pt2 -> finetune from .pt2 -> verify checkpoint."""
-        from deepmd.pt_expt.entrypoints.main import (
-            freeze,
-            main,
-        )
-
-        tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_pt2_")
-        old_cwd = os.getcwd()
-        os.chdir(tmpdir)
-        try:
-            # Phase 1: train pretrained model
-            config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1)
-            config = update_deepmd_input(config, warning=False)
-            config = normalize(config)
             ckpt_path = self._train_pretrained(config, tmpdir)
-
-            # Phase 2: freeze to .pt2
-            pt2_path = os.path.join(tmpdir, "frozen.pt2")
-            freeze(model=ckpt_path, output=pt2_path)
-            self.assertTrue(os.path.exists(pt2_path))
-
-            # Save pretrained weights before finetune overwrites model.ckpt.pt
-            pre_ckpt_copy = os.path.join(tmpdir, "pretrained_copy.pt")
-            shutil.copy2(ckpt_path, pre_ckpt_copy)
-
-            # Phase 3: finetune from .pt2 via CLI (lr=0 so weights stay unchanged)
-            ft_config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1)
-            ft_config["learning_rate"]["start_lr"] = 1e-30
-            ft_config["learning_rate"]["stop_lr"] = 1e-30
-            ft_config_file = os.path.join(tmpdir, "finetune_input.json")
-            with open(ft_config_file, "w") as f:
-                json.dump(ft_config, f)
-
-            main(
-                [
-                    "train",
-                    ft_config_file,
-                    "--finetune",
-                    pt2_path,
-                    "--skip-neighbor-stat",
-                ]
-            )
-
-            # Verify new checkpoint exists
-            ft_ckpt = os.path.join(tmpdir, "model.ckpt.pt")
-            self.assertTrue(os.path.exists(ft_ckpt), "Finetune checkpoint not found")
-
-            # Load finetuned model and verify it's valid
-            ft_state = torch.load(ft_ckpt, map_location=DEVICE, weights_only=True)
-            ft_model_state = ft_state["model"] if "model" in ft_state else ft_state
-            self.assertIn("_extra_state", ft_model_state)
-
-            # Load pretrained from saved copy (finetune overwrites model.ckpt.pt)
-            pre_state = torch.load(
-                pre_ckpt_copy, map_location=DEVICE, weights_only=True
-            )
-            pre_model_state = pre_state["model"] if "model" in pre_state else pre_state
-
-            # Inherited weights must match pretrained
-            self._assert_inherited_weights_match(
-                ft_model_state, pre_model_state, random_fitting=False
+            ft_config = normalize(
+                update_deepmd_input(
+                    _make_config(self.data_dir, target, numb_steps=1), warning=False
+                )
             )
-        finally:
-            os.chdir(old_cwd)
-            shutil.rmtree(tmpdir, ignore_errors=True)
-
-    def test_finetune_from_pt2_use_pretrain_script(self) -> None:
-        """Train -> freeze to .pt2 -> finetune with --use-pretrain-script."""
-        from deepmd.pt_expt.entrypoints.main import (
-            freeze,
-            main,
-        )
-
-        tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_pt2_ups_")
-        old_cwd = os.getcwd()
-        os.chdir(tmpdir)
-        try:
-            # Phase 1: train pretrained model
-            config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1)
-            config = update_deepmd_input(config, warning=False)
-            config = normalize(config)
-            ckpt_path = self._train_pretrained(config, tmpdir)
-
-            # Phase 2: freeze to .pt2 (embeds model_params)
-            pt2_path = os.path.join(tmpdir, "frozen.pt2")
-            freeze(model=ckpt_path, output=pt2_path)
-
-            # Phase 3: finetune from .pt2 with --use-pretrain-script
-            ft_model_params = deepcopy(model_se_e2_a)
-            ft_model_params["descriptor"]["neuron"] = [4, 8]  # different
-            ft_config = _make_config(self.data_dir, ft_model_params, numb_steps=1)
-            ft_config_file = os.path.join(tmpdir, "finetune_input.json")
-            with open(ft_config_file, "w") as f:
-                json.dump(ft_config, f)
-
-            main(
-                [
-                    "train",
-                    ft_config_file,
-                    "--finetune",
-                    pt2_path,
-                    "--use-pretrain-script",
-                    "--skip-neighbor-stat",
-                ]
+            ft_config["model"], finetune_links = get_finetune_rules(
+                ckpt_path, ft_config["model"]
             )
-
-            # Verify the output config was updated from pretrained
-            with open(os.path.join(tmpdir, "out.json")) as f:
-                output_config = json.load(f)
-            # Descriptor neuron should be from pretrained, not from ft_config
-            self.assertEqual(
-                output_config["model"]["descriptor"]["neuron"],
-                model_se_e2_a["descriptor"]["neuron"],
+            trainer_ft = get_trainer(
+                ft_config,
+                finetune_model=ckpt_path,
+                finetune_links=finetune_links,
             )
+            pretrained_state = torch.load(
+                ckpt_path, map_location=DEVICE, weights_only=True
+            )["model"]
+            ft_state = trainer_ft.wrapper.state_dict()
+            learned = [
+                key
+                for key in ft_state
+                if ".descriptor." in key
+                or (".fitting_net." in key and "bias_atom_e" not in key)
+            ]
+            self.assertTrue(learned)
+            for key in learned:
+                torch.testing.assert_close(ft_state[key], pretrained_state[key])
+            long_range = trainer_ft.wrapper.model["Default"].atomic_model.long_range
+            self.assertTrue(torch.all(long_range.k_grid_bounds > 0))
+            trainer_ft.run()
         finally:
             os.chdir(old_cwd)
             shutil.rmtree(tmpdir, ignore_errors=True)
diff --git a/source/tests/pt_expt/test_finetune_native_spin.py b/source/tests/pt_expt/test_finetune_native_spin.py
new file mode 100644
index 0000000000..88b453f36a
--- /dev/null
+++ b/source/tests/pt_expt/test_finetune_native_spin.py
@@ -0,0 +1,871 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Spin-free checkpoint transfer through both native-spin training backends."""
+
+import json
+from copy import (
+    deepcopy,
+)
+from importlib import (
+    import_module,
+)
+from pathlib import (
+    Path,
+)
+from typing import (
+    Any,
+)
+
+import numpy as np
+import pytest
+import torch
+
+from deepmd.pt.entrypoints.main import main as pt_main
+from deepmd.pt.model.model import get_model as get_pt_model
+from deepmd.pt.train.wrapper import ModelWrapper as PTModelWrapper
+from deepmd.pt.utils.finetune import (
+    _learned_part_path,
+)
+from deepmd.pt.utils.spin_finetune import (
+    compute_native_spin_input_stats,
+    get_native_spin_finetune_state,
+)
+from deepmd.pt_expt.entrypoints.main import (
+    main,
+)
+from deepmd.pt_expt.model import (
+    get_model,
+)
+from deepmd.pt_expt.train.wrapper import (
+    ModelWrapper,
+)
+from deepmd.pt_expt.utils.env import (
+    DEVICE,
+)
+from deepmd.utils.argcheck import (
+    normalize,
+)
+from deepmd.utils.bridging import (
+    expand_bridging_method,
+)
+from deepmd.utils.compat import (
+    update_deepmd_input,
+)
+from deepmd.utils.finetune import (
+    FinetuneRuleBuilder,
+)
+from deepmd.utils.model_preset import (
+    get_model_preset,
+)
+
+from .test_finetune import (
+    _make_config,
+    _native_spin_config,
+)
+
+_BACKENDS = [("pt", "dpa4"), ("pt_expt", "dpa4"), ("pt_expt", "dpa4c")]
+
+
+def _config(descriptor: str, native: bool, bridged: bool = False) -> dict:
+    config = _native_spin_config(descriptor, [True, False])
+    if not native:
+        config.pop("spin")
+    if bridged:
+        config["bridging_method"] = "zbl"
+    config["data_stat_nbatch"] = 1
+    return config
+
+
+def _canonical_config(config: dict) -> dict:
+    config = expand_bridging_method(config)
+    config["models"].reverse()
+    for child in config["models"]:
+        child.pop("type_map", None)
+    return config
+
+
+def _build(backend: str, config: dict) -> torch.nn.Module:
+    builder = get_pt_model if backend == "pt" else get_model
+    return builder(deepcopy(config)).to(DEVICE)
+
+
+def _learned(model: torch.nn.Module) -> torch.nn.Module:
+    return model.get_submodule(_learned_part_path(model).removesuffix("."))
+
+
+def _pretrained(backend: str, config: dict) -> torch.nn.Module:
+    """Build a geometric model whose nonzero forces make transfer observable."""
+    model = _build(backend, config)
+    generator = torch.Generator(device=DEVICE).manual_seed(31)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.add_(
+                0.02
+                * torch.randn(
+                    parameter.shape,
+                    dtype=parameter.dtype,
+                    device=parameter.device,
+                    generator=generator,
+                )
+            )
+        descriptor = _learned(model).descriptor
+        for name, start, stop in (("mean", -0.2, 0.2), ("stddev", 0.8, 1.2)):
+            array = getattr(descriptor, name)
+            array.copy_(
+                torch.linspace(
+                    start, stop, array.numel(), dtype=array.dtype, device=array.device
+                ).reshape(array.shape)
+            )
+    return model
+
+
+def _inputs() -> dict[str, torch.Tensor]:
+    return {
+        "coord": torch.tensor(
+            [[[0.0, 0.0, 0.0], [1.3, 0.2, 0.1], [0.3, 1.5, 0.4]]],
+            dtype=torch.float64,
+            device=DEVICE,
+        ),
+        "atype": torch.tensor([[0, 1, 0]], device=DEVICE),
+        "box": (6.0 * torch.eye(3, dtype=torch.float64, device=DEVICE)).reshape(1, 9),
+        "spin": torch.tensor(
+            [[[0.6, 0.2, -0.7], [0.1, 0.5, 0.1], [-0.4, 0.8, 0.3]]],
+            dtype=torch.float64,
+            device=DEVICE,
+        ),
+    }
+
+
+@pytest.mark.parametrize("backend, descriptor", _BACKENDS)
+@pytest.mark.parametrize("bridged", [False, True])
+def test_spin_free_transfer_preserves_function_and_trains_spin(
+    backend: str, descriptor: str, bridged: bool
+) -> None:
+    source = _pretrained(backend, _config(descriptor, False, bridged))
+    target = _build(backend, _config(descriptor, True, bridged))
+    target.load_state_dict(get_native_spin_finetune_state(target, source))
+    inputs = _inputs()
+    spin = inputs.pop("spin")
+    expected = source(**inputs, do_atomic_virial=True)
+    assert expected["force"].abs().max() > 1e-8
+    for scale in (0.0, 1.0, 2.0):
+        actual = target(**inputs, spin=scale * spin, do_atomic_virial=True)
+        for key in ("energy", "force", "virial", "atom_virial"):
+            torch.testing.assert_close(
+                actual[key], expected[key], rtol=1e-11, atol=1e-11
+            )
+        torch.testing.assert_close(
+            actual["force_mag"], torch.zeros_like(actual["force_mag"]), rtol=0, atol=0
+        )
+
+    spin_params = [param for name, param in target.named_parameters() if "spin" in name]
+    optimizer = torch.optim.SGD(spin_params, lr=1e-3)
+    (actual["force_mag"] - 0.2).square().sum().backward()
+    assert any(
+        param.grad is not None and param.grad.abs().max() > 1e-8
+        for param in spin_params
+    )
+    optimizer.step()
+    trained = target(**inputs, spin=spin)
+    assert trained["force_mag"].abs().max() > 1e-10
+    assert torch.count_nonzero(trained["force_mag"][:, 1]) == 0
+
+
+@pytest.mark.parametrize("backend, descriptor", _BACKENDS)
+@pytest.mark.parametrize("random_fitting", [False, True])
+def test_spin_free_transfer_preserves_target_bias_and_bridging(
+    backend: str, descriptor: str, random_fitting: bool
+) -> None:
+    source = _pretrained(backend, _config(descriptor, False))
+    target = _build(backend, _config(descriptor, True, bridged=True))
+    initial = {key: value.clone() for key, value in target.state_dict().items()}
+    with torch.no_grad():
+        for key, value in source.state_dict().items():
+            if key.rsplit(".", 1)[-1] in {"out_bias", "out_std", "bias_atom_e"}:
+                value.fill_(1234.0)
+    state = get_native_spin_finetune_state(
+        target, source, random_fitting=random_fitting
+    )
+    target.load_state_dict(state)
+    for key, value in state.items():
+        if key.rsplit(".", 1)[-1] in {"out_bias", "out_std", "bias_atom_e"} or (
+            random_fitting and ".fitting_net." in key
+        ):
+            torch.testing.assert_close(value, initial[key], rtol=0, atol=0)
+    source_descriptor = _learned(source).descriptor.state_dict()
+    target_descriptor = _learned(target).descriptor.state_dict()
+    for key, value in source_descriptor.items():
+        if value.shape == target_descriptor[key].shape:
+            torch.testing.assert_close(value, target_descriptor[key], rtol=0, atol=0)
+    np.testing.assert_array_equal(target.spin.use_spin, [True, False])
+
+
+@pytest.mark.parametrize("backend, descriptor", _BACKENDS)
+@pytest.mark.parametrize("width", [4, 16])
+def test_spin_free_transfer_rejects_unrelated_shape_change(
+    backend: str, descriptor: str, width: int
+) -> None:
+    source = _build(backend, _config(descriptor, False))
+    config = _config(descriptor, True)
+    config["fitting_net"]["neuron"] = [width]
+    target = _build(backend, config)
+    with pytest.raises(ValueError, match="Incompatible spin-free transfer"):
+        get_native_spin_finetune_state(target, source)
+
+
+@pytest.mark.parametrize(
+    "invalid, message",
+    [
+        ("source_without_learned_part", "one learned atomic model"),
+        ("target_without_learned_part", "one learned atomic model"),
+        ("type_map", "aligned type maps"),
+        ("descriptor_family", "same descriptor family"),
+        ("native_source", "spin-free source descriptor"),
+        ("spin_free_target", "native-spin target descriptor"),
+    ],
+)
+def test_spin_free_transfer_rejects_incompatible_models(
+    invalid: str, message: str
+) -> None:
+    source = _build("pt_expt", _config("dpa4", invalid == "native_source"))
+    target_config = _config(
+        "dpa4c" if invalid == "descriptor_family" else "dpa4",
+        invalid != "spin_free_target",
+    )
+    if invalid == "type_map":
+        target_config["type_map"] = ["C", "Fe"]
+    target = _build("pt_expt", target_config)
+    if invalid.endswith("without_learned_part"):
+        analytical = torch.nn.Module()
+        analytical.add_module("atomic_model", torch.nn.Module())
+        if invalid.startswith("source"):
+            source = analytical
+        else:
+            target = analytical
+    with pytest.raises(ValueError, match=message):
+        get_native_spin_finetune_state(target, source)
+
+
+@pytest.mark.parametrize("magnitude", [0.0, 2.0])
+def test_native_spin_statistics_preserve_pretrained_geometry(magnitude: float) -> None:
+    source = _pretrained("pt_expt", _config("dpa4c", False))
+    target = _build("pt_expt", _config("dpa4c", True))
+    target.load_state_dict(get_native_spin_finetune_state(target, source))
+    inputs = {
+        "coord": torch.tensor(
+            [[[0.0, 0.0, 0.0], [1.3, 0.0, 0.0]]], dtype=torch.float64, device=DEVICE
+        ),
+        "atype": torch.zeros((1, 2), dtype=torch.int64, device=DEVICE),
+        "spin": torch.tensor(
+            [[[0.0, 0.0, magnitude], [0.0, 0.0, magnitude]]],
+            dtype=torch.float64,
+            device=DEVICE,
+        ),
+    }
+    descriptor = _learned(target).descriptor
+    initial = {key: value.clone() for key, value in descriptor.state_dict().items()}
+    # Equivalent environments need no recalibration of pretrained geometry.
+    compute_native_spin_input_stats(target, lambda: [inputs])
+    torch.testing.assert_close(
+        descriptor.spin.spin_reference,
+        torch.tensor([magnitude or 1.0, 1.0, 1.0], dtype=torch.float64, device=DEVICE),
+        rtol=0,
+        atol=0,
+    )
+    for key, value in descriptor.state_dict().items():
+        if key != "spin.spin_reference":
+            torch.testing.assert_close(value, initial[key], rtol=0, atol=0)
+    actual = target(**inputs)
+    expected = source(coord=inputs["coord"], atype=inputs["atype"])
+    for key in ("energy", "force", "virial"):
+        torch.testing.assert_close(actual[key], expected[key], rtol=1e-11, atol=1e-11)
+    assert torch.count_nonzero(actual["force_mag"]) == 0
+
+
+def test_fresh_dpa4c_spin_statistics_keep_closed_gate_scale() -> None:
+    descriptor = _learned(_build("pt_expt", _config("dpa4c", True))).descriptor
+    descriptor.compute_input_stats([_inputs()])
+    start = descriptor.readout.get_dim_out()
+    stop = start + descriptor.spin.get_dim_out()
+    assert torch.count_nonzero(descriptor.spin.spin_gate) == 0
+    torch.testing.assert_close(
+        descriptor.mean[start:stop],
+        torch.zeros_like(descriptor.mean[start:stop]),
+        rtol=0,
+        atol=0,
+    )
+    torch.testing.assert_close(
+        descriptor.stddev[start:stop],
+        torch.ones_like(descriptor.stddev[start:stop]),
+        rtol=0,
+        atol=0,
+    )
+
+
+@pytest.mark.parametrize(
+    "fitting_type, neuron, source_form, target_form",
+    [
+        (fitting_type, neuron, "plain", "plain")
+        for fitting_type in ("dpa4_ener", "sezm_ener")
+        for neuron in ("omitted", None, [0], [0, 8], [8, 0])
+    ]
+    + [
+        ("dpa4_ener", "omitted", "concise", "canonical"),
+        ("sezm_ener", None, "concise", "canonical"),
+        ("sezm_ener", [0], "canonical", "concise"),
+        ("dpa4_ener", [0, 8], "canonical", "canonical"),
+    ],
+)
+def test_dpa4c_auto_width_inherits_source_architecture(
+    fitting_type: str,
+    neuron: str | list[int] | None,
+    source_form: str,
+    target_form: str,
+) -> None:
+    source_config = _config("dpa4c", False, bridged=source_form != "plain")
+    source_config["fitting_net"]["type"] = fitting_type
+    if neuron == "omitted":
+        source_config["fitting_net"].pop("neuron")
+    else:
+        source_config["fitting_net"]["neuron"] = neuron
+    target_config = deepcopy(source_config)
+    target_config["spin"] = {"scheme": "native", "use_spin": [True, False]}
+    if source_form == "canonical":
+        source_config = _canonical_config(source_config)
+    if target_form == "canonical":
+        target_config = _canonical_config(target_config)
+    original_source = deepcopy(source_config)
+    original = deepcopy(target_config)
+    source = _pretrained("pt_expt", source_config)
+    config, _ = FinetuneRuleBuilder(
+        source_config, target_config, change_model_params=False
+    ).build()
+    learned_config = config["models"][-1] if target_form == "canonical" else config
+    assert (
+        learned_config["fitting_net"]["neuron"] == _learned(source).fitting_net.neuron
+    )
+    assert target_config == original
+    assert source_config == original_source
+    target = _build("pt_expt", config)
+    target.load_state_dict(get_native_spin_finetune_state(target, source))
+    inputs = _inputs()
+    spin = inputs.pop("spin")
+    expected = source(**inputs)
+    actual = target(**inputs, spin=spin)
+    for key in ("energy", "force", "virial"):
+        torch.testing.assert_close(actual[key], expected[key], rtol=1e-11, atol=1e-11)
+
+
+@pytest.mark.parametrize("explicit_width", [False, True])
+@pytest.mark.parametrize("source_auto", [False, True])
+def test_dpa4c_width_aliases_preserve_explicit_config(
+    explicit_width: bool, source_auto: bool
+) -> None:
+    source_config = _config("dpa4c", False)
+    source_config["fitting_net"].update(
+        type="dpa4_ener", neuron=[0] if source_auto else [64]
+    )
+    source = _build("pt_expt", source_config)
+    source_script = deepcopy(source_config)
+    source_script["descriptor"]["type"] = "DPA4C"
+    source_script["fitting_net"]["n_neuron"] = source_script["fitting_net"].pop(
+        "neuron"
+    )
+    target = deepcopy(source_script)
+    target["spin"] = {"scheme": "native", "use_spin": [True, False]}
+    requested = [64] if explicit_width else [0]
+    target["fitting_net"]["n_neuron"] = requested
+    original = deepcopy(target)
+    resolved, _ = FinetuneRuleBuilder(
+        source_script, target, change_model_params=False
+    ).build()
+    config = normalize(
+        update_deepmd_input(_make_config("unused", resolved), warning=False)
+    )
+    model_config = config["model"]
+    expected = requested if explicit_width else _learned(source).fitting_net.neuron
+    assert model_config["fitting_net"]["neuron"] == expected
+    assert "n_neuron" not in model_config["fitting_net"]
+    assert target == original
+    model = _build("pt_expt", model_config)
+    if expected == _learned(source).fitting_net.neuron:
+        model.load_state_dict(get_native_spin_finetune_state(model, source))
+
+
+@pytest.mark.parametrize("resnet_dt", [False, True])
+@pytest.mark.parametrize("residual_side", ["source", "target"])
+@pytest.mark.parametrize("width_multiple", [1, 2])
+def test_dpa4c_transfer_rejects_changed_input_residual(
+    resnet_dt: bool, residual_side: str, width_multiple: int
+) -> None:
+    source_config = _config("dpa4c", False)
+    target_config = _config("dpa4c", True)
+    width_config = source_config if residual_side == "source" else target_config
+    width = _learned(_build("pt_expt", width_config)).descriptor.get_dim_out()
+    for config in (source_config, target_config):
+        config["fitting_net"].update(
+            neuron=[width_multiple * width], resnet_dt=resnet_dt
+        )
+    source = _pretrained("pt_expt", source_config)
+    target = _build("pt_expt", target_config)
+    initial = {key: value.clone() for key, value in target.state_dict().items()}
+    with pytest.raises(ValueError, match=r"residual.*input"):
+        get_native_spin_finetune_state(target, source)
+    for key, value in target.state_dict().items():
+        torch.testing.assert_close(value, initial[key], rtol=0, atol=0)
+
+    # A random fitting net does not inherit the source function or its residual.
+    state = get_native_spin_finetune_state(target, source, random_fitting=True)
+    target.load_state_dict(state)
+    for key, value in state.items():
+        if ".fitting_net." in key:
+            torch.testing.assert_close(value, initial[key], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("fitting_type", ["ener", "dpa4_ener", "sezm_ener"])
+@pytest.mark.parametrize("neuron", [[], [8], [8, 8]])
+def test_dpa4c_transfer_fitting_input_projection(
+    fitting_type: str, neuron: list[int]
+) -> None:
+    source_config = _config("dpa4c", False)
+    source_config["fitting_net"].update(type=fitting_type, neuron=neuron)
+    target_config = deepcopy(source_config)
+    target_config["spin"] = {"scheme": "native", "use_spin": [True, False]}
+    source = _pretrained("pt_expt", source_config)
+    target = _build("pt_expt", target_config)
+    # DPA4C's descriptor selects a shared fitting network at model construction.
+    assert _learned(source).fitting_net.nets.ndim == 0
+    assert _learned(target).fitting_net.nets.ndim == 0
+    state = get_native_spin_finetune_state(target, source)
+    assert all(not value.requires_grad for value in state.values())
+    target.load_state_dict(state)
+    inputs = _inputs()
+    spin = inputs.pop("spin")
+    expected = source(**inputs, do_atomic_virial=True)
+    assert expected["force"].abs().max() > 1e-8
+    for scale in (0.0, 1.0, 2.0):
+        actual = target(**inputs, spin=scale * spin, do_atomic_virial=True)
+        for key in ("energy", "force", "virial", "atom_virial"):
+            torch.testing.assert_close(
+                actual[key], expected[key], rtol=1e-11, atol=1e-11
+            )
+        assert torch.count_nonzero(actual["force_mag"]) == 0
+
+    (actual["force_mag"] - 0.2).square().sum().backward()
+    gate = _learned(target).descriptor.spin.spin_gate
+    assert gate.grad is not None and gate.grad.abs().max() > 1e-8
+    with torch.no_grad():
+        gate.add_(-1e-3 * gate.grad)
+    trained = target(**inputs, spin=spin)
+    assert trained["force_mag"].abs().max() > 1e-10
+    assert torch.count_nonzero(trained["force_mag"][:, 1]) == 0
+
+    restored = type(target).deserialize(target.serialize()).to(DEVICE)
+    reloaded = restored(**inputs, spin=spin)
+    for key in ("energy", "force", "virial", "force_mag"):
+        torch.testing.assert_close(reloaded[key], trained[key], rtol=1e-11, atol=1e-11)
+
+
+@pytest.mark.parametrize(
+    "backend, descriptor, source_type, target_type",
+    [
+        ("pt", "dpa4", "property", "dpa4_ener"),
+        ("pt_expt", "dpa4", "ener", "dpa4_ener"),
+        ("pt_expt", "dpa4c", "ener", "dpa4_ener"),
+        ("pt_expt", "dpa4c", "dpa4_ener", "ener"),
+    ],
+)
+def test_spin_free_transfer_rejects_changed_fitting_family(
+    backend: str, descriptor: str, source_type: str, target_type: str
+) -> None:
+    source_config = _config(descriptor, False)
+    target_config = _config(descriptor, True)
+    if backend == "pt_expt":
+        source_config["type"] = target_config["type"] = "standard"
+    for config, fitting_type in (
+        (source_config, source_type),
+        (target_config, target_type),
+    ):
+        config["fitting_net"].update(
+            type=fitting_type, neuron=[16 if fitting_type == "ener" else 8]
+        )
+        if fitting_type == "property":
+            config["fitting_net"].update(task_dim=1, property_name="energy")
+    source = _build(backend, source_config)
+    target = _build(backend, target_config)
+    initial = {key: value.clone() for key, value in target.state_dict().items()}
+    with pytest.raises(ValueError, match="matching fitting network types"):
+        get_native_spin_finetune_state(target, source)
+    for key, value in target.state_dict().items():
+        torch.testing.assert_close(value, initial[key], rtol=0, atol=0)
+    target.load_state_dict(
+        get_native_spin_finetune_state(target, source, random_fitting=True)
+    )
+    for key, value in target.state_dict().items():
+        if ".fitting_net." in key:
+            torch.testing.assert_close(value, initial[key], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    "backend, descriptor, source_neuron, target_neuron, trainable",
+    [
+        ("pt", "dpa4", [8], [8, 8], True),
+        ("pt", "dpa4", [8, 8], [8], True),
+        ("pt_expt", "dpa4", [8], [8, 8], True),
+        ("pt_expt", "dpa4c", [8], [8, 8], True),
+        ("pt_expt", "dpa4c", [8, 8], [8], True),
+        ("pt_expt", "dpa4c", [8], [8, 8], False),
+    ],
+)
+def test_spin_free_transfer_rejects_changed_fitting_layers(
+    backend: str,
+    descriptor: str,
+    source_neuron: list[int],
+    target_neuron: list[int],
+    trainable: bool,
+) -> None:
+    source_config = _config(descriptor, False)
+    target_config = _config(descriptor, True)
+    for config, neuron in (
+        (source_config, source_neuron),
+        (target_config, target_neuron),
+    ):
+        config["fitting_net"].update(
+            type="dpa4_ener", neuron=neuron, trainable=trainable
+        )
+    source = _build(backend, source_config)
+    target = _build(backend, target_config)
+    initial = {key: value.clone() for key, value in target.state_dict().items()}
+    with pytest.raises(ValueError, match="matching fitting network state keys"):
+        get_native_spin_finetune_state(target, source)
+    for key, value in target.state_dict().items():
+        torch.testing.assert_close(value, initial[key], rtol=0, atol=0)
+    target.load_state_dict(
+        get_native_spin_finetune_state(target, source, random_fitting=True)
+    )
+    for key, value in target.state_dict().items():
+        if ".fitting_net." in key:
+            torch.testing.assert_close(value, initial[key], rtol=0, atol=0)
+
+
+def test_spin_free_transfer_preserves_target_default_fparam() -> None:
+    source_config = _config("dpa4", False)
+    source_config["fitting_net"]["numb_fparam"] = 1
+    target_config = deepcopy(source_config)
+    target_config["spin"] = {"scheme": "native", "use_spin": [True, False]}
+    target_config["fitting_net"]["default_fparam"] = [0.25]
+    source = _build("pt", source_config)
+    target = _build("pt", target_config)
+    fitting = _learned(target).fitting_net
+    expected = fitting.default_fparam_tensor.clone()
+    target.load_state_dict(get_native_spin_finetune_state(target, source))
+    torch.testing.assert_close(fitting.default_fparam_tensor, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("fitting_type", ["ener", "dpa4_ener"])
+def test_dpa4c_spin_transfer_preserves_fitting_parameter_tail(
+    fitting_type: str,
+) -> None:
+    source_config = _config("dpa4c", False)
+    source_config["fitting_net"].update(type=fitting_type, numb_fparam=2, numb_aparam=1)
+    target_config = deepcopy(source_config)
+    target_config["spin"] = {"scheme": "native", "use_spin": [True, False]}
+    source = _pretrained("pt_expt", source_config)
+    target = _build("pt_expt", target_config)
+    target.load_state_dict(get_native_spin_finetune_state(target, source))
+    inputs = _inputs()
+    spin = inputs.pop("spin")
+    inputs["fparam"] = torch.tensor([[0.2, -0.1]], dtype=torch.float64, device=DEVICE)
+    inputs["aparam"] = torch.tensor(
+        [[[0.3], [0.5], [0.7]]], dtype=torch.float64, device=DEVICE
+    )
+    expected = source(**inputs)
+    actual = target(**inputs, spin=spin)
+    for key in ("energy", "force", "virial"):
+        torch.testing.assert_close(actual[key], expected[key], rtol=1e-11, atol=1e-11)
+
+
+@pytest.mark.parametrize("grade", ["nano", "mini", "neo", "air", "plus"])
+def test_dpa4c_preset_spin_free_transfer(grade: str) -> None:
+    source_config = get_model_preset(f"dpa4c-{grade}-v20260911")
+    source_config["descriptor"]["seed"] = 17
+    source_config["fitting_net"]["seed"] = 19
+    target_config = deepcopy(source_config)
+    target_config["spin"] = {"scheme": "native", "use_spin": ["Fe"]}
+    source = _pretrained("pt_expt", source_config)
+    target = _build("pt_expt", target_config)
+    target.load_state_dict(get_native_spin_finetune_state(target, source))
+    inputs = _inputs()
+    spin = inputs.pop("spin")
+    inputs["atype"] = torch.tensor(
+        [[source.get_type_map().index(name) for name in ("Fe", "C", "Fe")]],
+        device=DEVICE,
+    )
+    expected = source(**inputs)
+    actual = target(**inputs, spin=spin)
+    for key in ("energy", "force", "virial"):
+        torch.testing.assert_close(actual[key], expected[key], rtol=2e-5, atol=2e-5)
+    assert torch.count_nonzero(actual["force_mag"]) == 0
+
+
+@pytest.mark.parametrize(
+    "backend, descriptor, fitting_type",
+    [
+        ("pt", "dpa4", "dpa4_ener"),
+        ("pt_expt", "dpa4", "dpa4_ener"),
+        ("pt_expt", "dpa4c", "ener"),
+        ("pt_expt", "dpa4c", "dpa4_ener"),
+    ],
+)
+@pytest.mark.parametrize("bridged", [False, True])
+@pytest.mark.parametrize("multitask", [False, True])
+def test_spin_free_native_finetune_cli(
+    backend: str,
+    descriptor: str,
+    fitting_type: str,
+    bridged: bool,
+    multitask: bool,
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _run_spin_free_native_finetune_cli(
+        backend,
+        descriptor,
+        fitting_type,
+        bridged,
+        multitask,
+        "transfer",
+        tmp_path,
+        monkeypatch,
+    )
+
+
+@pytest.mark.parametrize("backend, descriptor", [("pt", "dpa4"), ("pt_expt", "dpa4c")])
+@pytest.mark.parametrize("mode", ["random", "pretrain_script"])
+@pytest.mark.parametrize("multitask", [False, True])
+def test_spin_free_native_finetune_cli_options(
+    backend: str,
+    descriptor: str,
+    mode: str,
+    multitask: bool,
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _run_spin_free_native_finetune_cli(
+        backend, descriptor, "dpa4_ener", False, multitask, mode, tmp_path, monkeypatch
+    )
+
+
+@pytest.mark.parametrize("mode", ["transfer", "random", "pretrain_script"])
+@pytest.mark.parametrize("multitask", [False, True])
+def test_dpa4c_auto_width_finetune_cli(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mode: str, multitask: bool
+) -> None:
+    _run_spin_free_native_finetune_cli(
+        "pt_expt",
+        "dpa4c",
+        "dpa4_ener",
+        False,
+        multitask,
+        mode,
+        tmp_path,
+        monkeypatch,
+        auto_width=True,
+    )
+
+
+def test_dpa4c_canonical_auto_width_finetune_cli(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    _run_spin_free_native_finetune_cli(
+        "pt_expt",
+        "dpa4c",
+        "dpa4_ener",
+        True,
+        False,
+        "transfer",
+        tmp_path,
+        monkeypatch,
+        auto_width=True,
+        canonical=True,
+    )
+
+
+def _run_spin_free_native_finetune_cli(
+    backend: str,
+    descriptor: str,
+    fitting_type: str,
+    bridged: bool,
+    multitask: bool,
+    mode: str,
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    auto_width: bool = False,
+    canonical: bool = False,
+) -> None:
+    source_config = _config(descriptor, False)
+    source_config["fitting_net"]["type"] = fitting_type
+    if auto_width:
+        source_config["fitting_net"].pop("neuron")
+    source = _pretrained(backend, source_config)
+    wrapper_cls = PTModelWrapper if backend == "pt" else ModelWrapper
+    wrapper = wrapper_cls(
+        {"magnetic": source} if multitask else source,
+        model_params={"model_dict": {"magnetic": source_config}}
+        if multitask
+        else source_config,
+    )
+    wrapper.train_infos["step"] = 100
+    checkpoint = tmp_path / "pretrained.pt"
+    torch.save(
+        {
+            "model": wrapper.state_dict(),
+            "optimizer": {"source_only": True},
+            "ema": {"source_only": True},
+        },
+        checkpoint,
+    )
+
+    inputs = _inputs()
+    spin = inputs.pop("spin")
+    expected = source(**inputs)
+    system = tmp_path / "data" / "data_0"
+    arrays = system / "set.000"
+    arrays.mkdir(parents=True)
+    (system / "type.raw").write_text("0\n1\n0\n")
+    (system / "type_map.raw").write_text("Fe\nC\n")
+    labels = {
+        "coord": inputs["coord"],
+        "box": inputs["box"],
+        "spin": spin,
+        "energy": expected["energy"] + 1.75,
+        "force": expected["force"],
+        "force_mag": torch.full_like(spin, 0.2),
+    }
+    for key, value in labels.items():
+        data = value.detach().cpu().numpy().reshape(1, -1)
+        np.save(arrays / f"{key}.npy", np.repeat(data, 2, axis=0))
+
+    target_config = _config(descriptor, True, bridged)
+    target_config["fitting_net"]["type"] = fitting_type
+    if auto_width:
+        target_config["fitting_net"].pop("neuron")
+    options = []
+    if mode == "random":
+        initial_fitting = dict(
+            _learned(_build(backend, target_config)).fitting_net.named_parameters()
+        )
+        if multitask:
+            target_config["finetune_head"] = "RANDOM"
+        else:
+            options = ["--model-branch", "RANDOM"]
+    elif mode == "pretrain_script":
+        target_config["descriptor"]["channels"] *= 2
+        options = ["--use-pretrain-script"]
+    if bridged:
+        target_config["preset_out_bias"] = {"energy": {"Fe": 1.25, "C": -0.75}}
+    if canonical:
+        target_config = _canonical_config(target_config)
+    # A zero-step run exposes RANDOM initialization before optimization.
+    steps = 0 if mode == "random" else 2
+    config = _make_config(str(system.parent), target_config, numb_steps=steps)
+    config["loss"] = {
+        "type": "ener_spin",
+        "start_pref_e": 1.0,
+        "limit_pref_e": 1.0,
+        "start_pref_fr": 1.0,
+        "limit_pref_fr": 1.0,
+        "start_pref_fm": 1.0,
+        "limit_pref_fm": 1.0,
+    }
+    config["training"].update(enable_ema=True, ema_decay=0.9)
+    if multitask:
+        config["model"] = {"model_dict": {"magnetic": target_config}}
+        config["loss_dict"] = {"magnetic": config.pop("loss")}
+        config["training"]["data_dict"] = {
+            "magnetic": {
+                key: config["training"].pop(key)
+                for key in ("training_data", "validation_data")
+            }
+        }
+    input_path = tmp_path / "input.json"
+    input_path.write_text(json.dumps(config))
+    monkeypatch.chdir(tmp_path)
+    entrypoint = pt_main if backend == "pt" else main
+    trainers = []
+    if mode == "random":
+        entrypoints = import_module(f"deepmd.{backend}.entrypoints.main")
+        get_trainer = entrypoints.get_trainer
+
+        def capture_trainer(*args: Any, **kwargs: Any) -> Any:
+            trainer = get_trainer(*args, **kwargs)
+            trainers.append(trainer)
+            return trainer
+
+        monkeypatch.setattr(entrypoints, "get_trainer", capture_trainer)
+    entrypoint(
+        [
+            "train",
+            str(input_path),
+            "--finetune",
+            str(checkpoint),
+            "--skip-neighbor-stat",
+            *options,
+        ]
+    )
+    if mode == "random":
+        assert len(trainers) == 1
+        model = trainers[0].model
+        trained = model["magnetic"] if multitask else model
+        for key, value in _learned(trained).fitting_net.named_parameters():
+            torch.testing.assert_close(value, initial_fitting[key], rtol=0, atol=0)
+        assert torch.count_nonzero(trained(**inputs, spin=spin)["force_mag"]) == 0
+        return
+    saved = torch.load(
+        tmp_path / "model.ckpt.pt", map_location=DEVICE, weights_only=True
+    )
+    params = saved["model"]["_extra_state"]["model_params"]
+    model_params = params["model_dict"]["magnetic"] if multitask else params
+    learned_params = model_params["models"][-1] if canonical else model_params
+    if auto_width:
+        assert (
+            learned_params["fitting_net"]["neuron"]
+            == _learned(source).fitting_net.neuron
+        )
+    assert model_params["spin"]["scheme"] == "native"
+    assert model_params["spin"]["use_spin"] == [True, False]
+    assert (
+        learned_params["descriptor"]["channels"]
+        == source_config["descriptor"]["channels"]
+    )
+    # PT checkpoints record the zero-based step; pt_expt records the step count.
+    assert saved["model"]["_extra_state"]["train_infos"]["step"] == (
+        max(steps - 1, 0) if backend == "pt" else steps
+    )
+    assert "source_only" not in saved["optimizer"]
+    assert "source_only" not in saved["ema"]
+    trained = _build(backend, model_params)
+    loaded = wrapper_cls(
+        {"magnetic": trained} if multitask else trained, model_params=params
+    )
+    loaded.load_state_dict(saved["model"])
+    if descriptor == "dpa4c":
+        reference = _learned(trained).descriptor.spin.spin_reference
+        expected_reference = spin[:, [0, 2]].square().sum(dim=-1).mean().sqrt()
+        torch.testing.assert_close(reference[0], expected_reference, rtol=1e-12, atol=0)
+    bias = trained.get_out_bias().reshape(-1)
+    if bridged:
+        torch.testing.assert_close(
+            bias,
+            torch.tensor([1.25, -0.75], dtype=torch.float64, device=DEVICE),
+            rtol=0,
+            atol=0,
+        )
+    else:
+        torch.testing.assert_close(
+            2.0 * bias[0] + bias[1],
+            torch.tensor(1.75, dtype=torch.float64, device=DEVICE),
+            rtol=1e-11,
+            atol=1e-11,
+        )
+    assert trained(**inputs, spin=spin)["force_mag"].abs().max() > 1e-10
diff --git a/source/tests/pt_expt/test_freeze_compress.py b/source/tests/pt_expt/test_freeze_compress.py
new file mode 100644
index 0000000000..d1a9407b22
--- /dev/null
+++ b/source/tests/pt_expt/test_freeze_compress.py
@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Single-pass compressed export through the CLI and ASE calculator."""
+
+import json
+import logging
+import zipfile
+from pathlib import (
+    Path,
+)
+from unittest import (
+    mock,
+)
+
+import numpy as np
+import pytest
+import torch
+from ase import (
+    Atoms,
+)
+
+from deepmd.calculator import (
+    DP,
+)
+from deepmd.infer import (
+    DeepPot,
+)
+from deepmd.main import (
+    parse_args,
+)
+from deepmd.pt_expt.entrypoints import compress as compression
+from deepmd.pt_expt.entrypoints import main as entrypoints
+from deepmd.pt_expt.model.get_model import (
+    get_model,
+)
+from deepmd.pt_expt.train.wrapper import (
+    ModelWrapper,
+)
+from deepmd.pt_expt.utils import (
+    env,
+    serialization,
+)
+from deepmd.utils.vacuum_reference import (
+    reference_charge_spin,
+)
+
+from .test_calculator import (
+    _config,
+    _write_checkpoint,
+)
+
+
+@pytest.fixture(autouse=True)
+def cpu_export(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(env, "DEVICE", torch.device("cpu"))
+    for name in ("DP_CUDA_INFER", "DP_TRITON_INFER", "DP_COMPILE_INFER"):
+        monkeypatch.delenv(name, raising=False)
+
+
+@pytest.mark.parametrize("compress", [False, True])
+def test_freeze_cli_compression_option(tmp_path: Path, compress: bool) -> None:
+    checkpoint = tmp_path / "model.pt"
+    checkpoint.touch()
+    args = ["--pt-expt", "freeze", "-c", str(checkpoint)]
+    if compress:
+        args.append("--compress")
+    flags = parse_args(args)
+    with mock.patch.object(entrypoints, "freeze") as freeze:
+        entrypoints.main(flags)
+    assert freeze.call_args.kwargs["compress"] is compress
+
+
+@pytest.mark.parametrize("backend", ["pt", "tf", "pd", "jax"])
+def test_freeze_cli_rejects_unsupported_compression(
+    backend: str, capsys: pytest.CaptureFixture[str]
+) -> None:
+    """Unsupported backends reject compression instead of ignoring the flag."""
+    with pytest.raises(SystemExit) as error:
+        parse_args([f"--{backend}", "freeze", "--compress"])
+    assert error.value.code == 2
+    assert "freeze --compress requires --pt-expt" in capsys.readouterr().err
+
+
+def test_freeze_missing_neighbor_statistics_explains_compression_command(
+    tmp_path: Path,
+) -> None:
+    """Models needing data statistics point to the CLI that accepts them."""
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", "dense")
+    output = tmp_path / "model.pte"
+    with pytest.raises(ValueError, match="dp --pt-expt compress") as error:
+        entrypoints.freeze(str(checkpoint), str(output), compress=True)
+    assert "without --compress" in str(error.value)
+    assert not output.exists()
+
+
+@pytest.mark.parametrize("compress", [False, True])
+def test_freeze_compression_changes_snapshot_before_export(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, compress: bool
+) -> None:
+    config = _config("dpa4c")
+    config["descriptor"].update(add_chg_spin_ebd=True, default_chg_spin=[1, 2])
+    config["fitting_net"]["vacuum_ref"] = True
+    config["preset_out_bias"] = {"energy": {"O": 0.0, "H": 0.0}}
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4c", config=config)
+    original = checkpoint.read_bytes()
+    reference = get_model(config).eval()
+    ModelWrapper(reference, model_params=config).load_state_dict(
+        torch.load(checkpoint, map_location="cpu", weights_only=True)["model"]
+    )
+    fitting = reference.atomic_model.fitting_net
+    original_bias = fitting.bias_atom_e.detach().numpy().copy()
+    fitting.vacuum_ref = False
+    with torch.no_grad():
+        fitting.bias_atom_e.zero_()
+        reference.atomic_model.out_bias.zero_()
+    coord = torch.zeros((2, 1, 3), dtype=torch.float64)
+    atype = torch.arange(2, dtype=torch.int64).reshape(2, 1)
+    neutral = (
+        reference.call_common(
+            coord,
+            atype,
+            None,
+            charge_spin=torch.as_tensor(reference_charge_spin(config["type_map"])),
+        )["energy"]
+        .detach()
+        .numpy()
+        .reshape(2, 1)
+    )
+    deployment = (
+        reference.call_common(coord, atype, None)["energy"]
+        .detach()
+        .numpy()
+        .reshape(2, 1)
+    )
+    # The deployment state must differ from the neutral isolated-atom reference.
+    assert np.max(np.abs(deployment - neutral)) > 1e-5
+    enable_compression = type(reference).enable_compression
+
+    def check_reference_before_compression(
+        model: torch.nn.Module, *args: object, **kwargs: object
+    ) -> None:
+        np.testing.assert_allclose(
+            model.atomic_model.fitting_net.bias_atom_e.detach().numpy(),
+            original_bias - neutral,
+            rtol=1e-6,
+            atol=1e-7,
+        )
+        enable_compression(model, *args, **kwargs)
+
+    monkeypatch.setattr(
+        type(reference), "enable_compression", check_reference_before_compression
+    )
+    captured = {}
+
+    def capture(model_file: str, data: dict, **kwargs: object) -> None:
+        captured.update(data=data, options=kwargs)
+
+    monkeypatch.setattr(serialization, "deserialize_to_file", capture)
+    monkeypatch.setattr(compression, "deserialize_to_file", capture)
+    entrypoints.freeze(
+        str(checkpoint),
+        str(tmp_path / "model.pt2"),
+        compress=compress,
+        with_comm=False,
+    )
+    model = captured["data"]["model"]
+    assert ("compress" in model["descriptor"]) is compress
+    assert model["fitting"]["vacuum_ref"] is (not compress)
+    np.testing.assert_allclose(
+        model["fitting"]["@variables"]["bias_atom_e"],
+        original_bias - neutral if compress else original_bias,
+        rtol=1e-6,
+        atol=1e-7,
+    )
+    assert captured["options"]["with_comm"] is False
+    assert captured["options"]["lower_kind"] == ("auto" if compress else "graph")
+    assert checkpoint.read_bytes() == original
+
+
+@pytest.mark.parametrize("case", ["precision", "exclusions", "charge_state", "widths"])
+def test_ase_falls_back_but_explicit_compression_reports_ineligibility(
+    tmp_path: Path, caplog: pytest.LogCaptureFixture, case: str
+) -> None:
+    config = _config("dpa4c")
+    if case == "precision":
+        config["descriptor"]["precision"] = "float64"
+        config["fitting_net"]["precision"] = "float64"
+        error = "requires descriptor precision"
+    elif case == "exclusions":
+        config["descriptor"]["exclude_types"] = [[0, 1]]
+        error = "non-empty `exclude_types`"
+    elif case == "widths":
+        config["descriptor"]["radial_modes"] = 3
+        error = "radial_modes"
+    else:
+        config["descriptor"]["add_chg_spin_ebd"] = True
+        config["descriptor"]["default_chg_spin"] = None
+        error = "requires `default_chg_spin`"
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4c", config=config)
+    loaded = DeepPot(str(checkpoint))
+
+    def write_artifact(source: str, output: str, **kwargs: object) -> None:
+        Path(output).write_bytes(b"export")
+
+    with (
+        mock.patch("deepmd.calculator.DeepPot", return_value=loaded),
+        mock.patch.object(entrypoints, "freeze", side_effect=write_artifact) as freeze,
+        caplog.at_level(logging.WARNING, logger="deepmd.calculator"),
+    ):
+        DP(checkpoint, freeze=True)
+    assert freeze.call_args.kwargs["compress"] is False
+    assert "Freezing without compression" in caplog.text
+    assert error in caplog.text
+    output = tmp_path / "explicit.pt2"
+    with pytest.raises(ValueError, match=error):
+        entrypoints.freeze(str(checkpoint), str(output), compress=True)
+    assert not output.exists()
+
+
+def test_compressed_freeze_compiles_one_cpu_artifact(tmp_path: Path) -> None:
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4c")
+    output = tmp_path / "model.pt2"
+
+    def compile_package(
+        exported: torch.export.ExportedProgram, package_path: str, **kwargs: object
+    ) -> None:
+        with zipfile.ZipFile(package_path, "w") as archive:
+            archive.writestr("model/data.pkl", b"")
+
+    with mock.patch(
+        "torch._inductor.aoti_compile_and_package", side_effect=compile_package
+    ) as compile_mock:
+        entrypoints.freeze(str(checkpoint), str(output), compress=True, with_comm=False)
+    assert compile_mock.call_count == 1
+    with zipfile.ZipFile(output) as archive:
+        metadata = json.loads(archive.read("model/extra/metadata.json"))
+        assert metadata["has_comm_artifact"] is False
+        assert metadata["lower_input_kind"] == "graph"
+        assert "model/extra/forward_lower_with_comm.pt2" not in archive.namelist()
+    snapshot = serialization.serialize_from_file(str(output))
+    assert "compress" in snapshot["model"]["descriptor"]
+
+
+def test_ase_float64_fallback_exports_and_evaluates(
+    tmp_path: Path, caplog: pytest.LogCaptureFixture
+) -> None:
+    config = _config("dpa4c")
+    config["descriptor"]["precision"] = "float64"
+    config["fitting_net"]["precision"] = "float64"
+    checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4c", config=config)
+    atoms = Atoms(
+        "OH2",
+        positions=[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [-0.3, 0.8, 0.0]],
+        cell=[8.0, 8.0, 8.0],
+        pbc=True,
+        calculator=DP(checkpoint),
+    )
+    expected = atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress()
+    caplog.clear()
+    with caplog.at_level(logging.WARNING, logger="deepmd.calculator"):
+        atoms.calc = DP(checkpoint, freeze=True)
+    actual = atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress()
+    assert "Freezing without compression" in caplog.text
+    snapshot = serialization.serialize_from_file(str(tmp_path / "model-ase_single.pt2"))
+    assert "compress" not in snapshot["model"]["descriptor"]
+    for reference, result in zip(expected, actual, strict=True):
+        np.testing.assert_allclose(result, reference, rtol=1e-6, atol=1e-8)
diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py
index 539b72f2c8..0a23567a22 100644
--- a/source/tests/pt_expt/test_lmdb_training.py
+++ b/source/tests/pt_expt/test_lmdb_training.py
@@ -22,6 +22,7 @@
 import msgpack
 import numpy as np
 
+import deepmd.utils.random as dp_random
 from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module
 from deepmd.dpmodel.utils.batch import (
     normalize_batch,
@@ -35,6 +36,7 @@
     collate_lmdb_frames,
 )
 from deepmd.pt_expt.entrypoints.main import (
+    _build_data_system,
     get_trainer,
 )
 from deepmd.pt_expt.loss import (
@@ -43,6 +45,9 @@
 from deepmd.pt_expt.train.training import (
     Trainer,
 )
+from deepmd.pt_expt.train.validation import (
+    FullValidator,
+)
 from deepmd.pt_expt.utils.lmdb_dataset import (
     LmdbDataSystem,
 )
@@ -58,6 +63,12 @@
 from deepmd.utils.data import (
     DataRequirementItem,
 )
+from deepmd.utils.data_group import (
+    GroupedDataSystem,
+)
+from deepmd.utils.data_system import (
+    get_neighbor_stat_data,
+)
 
 from .compile_utils import (
     REQUIRES_SUPPORTED_COMPILE,
@@ -452,7 +463,7 @@ def test_ragged_filter_uses_the_consumed_pass_length(self) -> None:
                 trainer.multi_task = False
                 trainer.world_size = 1
                 trainer.training_data_by_task = {"Default": data}
-                trainer.min_pair_dist_by_task = {"Default": 0.5}
+                trainer.pair_filter_enabled_by_task = {"Default": True}
                 try:
                     # This shuffle packs five batches and leaves frame 3 last;
                     # the following shuffle packs four batches.
@@ -700,6 +711,115 @@ def test_stat_batches_restart_after_non_divisible_pass(self) -> None:
             )
 
 
+class TestLmdbDatasetGroups(unittest.TestCase):
+    def setUp(self) -> None:
+        directory = tempfile.TemporaryDirectory()
+        self.addCleanup(directory.cleanup)
+        self.directory = directory.name
+        self.first = os.path.join(self.directory, "first.lmdb")
+        self.parent = os.path.join(self.directory, "nested")
+        os.mkdir(self.parent)
+        self.second = os.path.join(self.parent, "second.lmdb")
+        _create_test_lmdb(self.first, nframes=2, natoms=6)
+        _create_test_lmdb(self.second, nframes=8, natoms=8)
+        workers = patch.dict(os.environ, {"DP_LMDB_NUM_WORKERS": "0"})
+        workers.start()
+        self.addCleanup(workers.stop)
+
+    def test_unseeded_distributed_groups_share_source_choices(self) -> None:
+        sequences = []
+        for rank in range(2):
+            dp_random.seed(17 + rank)
+            data = _build_data_system(
+                {
+                    "systems": [self.first, self.second],
+                    "batch_size": 1,
+                    "auto_prob": "prob_uniform",
+                },
+                ["O", "H"],
+                rank=rank,
+                world_size=2,
+            )
+            self.addCleanup(data.close)
+            sequences.append([data.get_batch()["coord"].shape[1] for _ in range(24)])
+        self.assertEqual(sequences[0], sequences[1])
+
+    def test_multiple_paths_and_recursive_parent_use_flat_sampling(self) -> None:
+        for systems in ([self.first, self.parent], self.directory):
+            with self.subTest(systems=systems):
+                data = _build_data_system(
+                    {"systems": systems, "batch_size": 1, "auto_prob": "prob_sys_size"},
+                    ["O", "H"],
+                    seed=0,
+                )
+                self.addCleanup(data.close)
+                self.assertIsInstance(data, GroupedDataSystem)
+                np.testing.assert_allclose(data.sys_probs, [0.2, 0.8])
+                shapes = {data.get_batch()["coord"].shape[1] for _ in range(20)}
+                self.assertEqual(shapes, {6, 8})
+
+    def test_group_probability_and_statistics_keep_sources_separate(self) -> None:
+        data = _build_data_system(
+            [
+                {"systems": self.first, "batch_size": 1, "prob": 0.6},
+                {"systems": self.parent, "batch_size": 2, "prob": 0.4},
+            ],
+            ["O", "H"],
+            seed=0,
+        )
+        self.addCleanup(data.close)
+        np.testing.assert_allclose(data.sys_probs, [0.6, 0.4])
+        samples = make_stat_input(data, 1)
+        self.assertEqual(len(samples), 2)
+        self.assertEqual([sample["coord"].shape[1] for sample in samples], [6, 8])
+        self.assertEqual([sample["coord"].shape[0] for sample in samples], [1, 2])
+
+    def test_full_validation_caches_each_lmdb_separately(self) -> None:
+        data = _build_data_system(
+            {"systems": [self.first, self.second], "batch_size": 1}, ["O", "H"]
+        )
+        self.addCleanup(data.close)
+        validator = FullValidator.__new__(FullValidator)
+        validator.validation_data = data
+        validator._lmdb_snapshots = {}
+        first = list(validator._iter_validation_data_systems())
+        snapshots = dict(validator._lmdb_snapshots)
+        second = list(validator._iter_validation_data_systems())
+        self.assertEqual(len(first), 2)
+        self.assertEqual(len(snapshots), 2)
+        self.assertEqual(validator._lmdb_snapshots, snapshots)
+        self.assertEqual(
+            [system.get_test()["coord"].shape[0] for system in second], [2, 8]
+        )
+
+    def test_distributed_epoch_accounts_for_each_leaf_padding(self) -> None:
+        paths = [
+            os.path.join(self.directory, f"tiny_{index}.lmdb") for index in range(2)
+        ]
+        for path in paths:
+            _create_test_lmdb(path, nframes=1, natoms=6)
+        for rank in (0, 1):
+            data = _build_data_system(
+                {"systems": paths, "batch_size": 1},
+                ["O", "H"],
+                seed=0,
+                rank=rank,
+                world_size=2,
+            )
+            self.addCleanup(data.close)
+            trainer = Trainer.__new__(Trainer)
+            trainer.training_data_by_task = {"Default": data}
+            trainer.world_size = 2
+            self.assertEqual(trainer._epoch_length("Default"), 2)
+
+    def test_neighbor_statistics_visit_both_lmdbs(self) -> None:
+        data = get_neighbor_stat_data(
+            {"systems": [self.first, self.parent], "batch_size": 1}, ["O", "H"]
+        )
+        self.assertEqual(data.get_ntypes(), 2)
+        self.assertEqual([system.get_natoms() for system in data.data_systems], [6, 8])
+
+
 class TestLmdbTrainingLoop(unittest.TestCase):
     """End-to-end: get_trainer routes an LMDB path and runs training steps."""
 
@@ -786,6 +906,31 @@ def test_get_trainer_routes_lmdb(self) -> None:
         finally:
             os.chdir(cwd)
 
+    def test_get_trainer_runs_weighted_lmdb_groups(self) -> None:
+        config = self._make_lmdb_config(numb_steps=2)
+        config["training"]["training_data"] = [
+            {
+                "systems": self.lmdb_path,
+                "batch_size": 1,
+                "prob": 0.6,
+                "min_pair_dist": 0.1,
+            },
+            {"systems": self.val_lmdb_path, "batch_size": 2, "prob": 0.4},
+        ]
+        config["training"]["validation_data"]["systems"] = [
+            self.lmdb_path,
+            self.val_lmdb_path,
+        ]
+        config = normalize(update_deepmd_input(config, warning=False))
+        cwd = os.getcwd()
+        os.chdir(self.tmpdir)
+        try:
+            trainer = get_trainer(config)
+            np.testing.assert_allclose(trainer.training_data.sys_probs, [0.6, 0.4])
+            trainer.run()
+        finally:
+            os.chdir(cwd)
+
     def test_numb_epoch_counts_passes_over_the_lmdb(self) -> None:
         """One epoch is one pass over the frames of the LMDB."""
         config = self._make_lmdb_config()
@@ -914,14 +1059,12 @@ class TestRaggedTrainingBatches(unittest.TestCase):
 
     def setUp(self) -> None:
         self.tmpdir = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.tmpdir)
         self.lmdb_path = os.path.join(self.tmpdir, "mixed.lmdb")
         self.spin_lmdb_path = os.path.join(self.tmpdir, "mixed-spin.lmdb")
         _create_mixed_nloc_test_lmdb(self.lmdb_path)
         _create_mixed_nloc_test_lmdb(self.spin_lmdb_path, include_spin=True)
 
-    def tearDown(self) -> None:
-        shutil.rmtree(self.tmpdir, ignore_errors=True)
-
     def _config_for_model(
         self,
         model: dict,
@@ -1064,6 +1207,7 @@ def _run(self, descriptor: dict, *, compile: bool = False):
         os.chdir(self.tmpdir)
         try:
             trainer = get_trainer(config)
+            self.addCleanup(trainer._close_data_systems)
             batch = trainer.training_data.get_batch()
             trainer.run()
             return trainer, batch
@@ -1081,6 +1225,7 @@ def _run_spin(self, model: dict):
         os.chdir(self.tmpdir)
         try:
             trainer = get_trainer(config)
+            self.addCleanup(trainer._close_data_systems)
             batch = trainer.training_data.get_batch()
             trainer.run()
             return trainer, batch
@@ -1098,7 +1243,9 @@ def test_min_pair_dist_filters_ragged_batches(self) -> None:
             # rather than its minimum, so the distances that set a splitting
             # threshold are measured on the geometry itself. A unit pair
             # threshold makes each margin the minimum pair distance in Å.
-            batch = get_trainer(copy.deepcopy(config)).training_data.get_batch()
+            probe = get_trainer(copy.deepcopy(config))
+            self.addCleanup(probe._close_data_systems)
+            batch = probe.training_data.get_batch()
             distances = compute_min_pair_margin_batch(
                 np.asarray(batch["coord"]),
                 np.asarray(batch["box"]),
@@ -1111,6 +1258,7 @@ def test_min_pair_dist_filters_ragged_batches(self) -> None:
 
             config["training"]["training_data"]["min_pair_dist"] = threshold
             trainer = get_trainer(copy.deepcopy(config))
+            self.addCleanup(trainer._close_data_systems)
             self.assertTrue(trainer.training_data._reader.ragged_batches)
             dropped = False
             for _ in range(6):
@@ -1163,6 +1311,7 @@ def test_loss_contract_can_require_rectangular_batches(self) -> None:
                     config = self._config(self._dpa1())
                     config["loss"].update(loss_params)
                     trainer = get_trainer(config)
+                    self.addCleanup(trainer._close_data_systems)
                     self.assertFalse(trainer.training_data._reader.ragged_batches)
         finally:
             os.chdir(cwd)
@@ -1175,6 +1324,7 @@ def test_non_mixing_rule_keeps_the_rectangular_layout(self) -> None:
         os.chdir(self.tmpdir)
         try:
             trainer = get_trainer(config)
+            self.addCleanup(trainer._close_data_systems)
             batch = trainer.training_data.get_batch()
         finally:
             os.chdir(cwd)
@@ -1259,6 +1409,7 @@ def record_layout(trainer_self, model_key: str) -> int:
         try:
             with patch.object(Trainer, "_epoch_length", record_layout):
                 trainer = get_trainer(config)
+                self.addCleanup(trainer._close_data_systems)
         finally:
             os.chdir(cwd)
 
diff --git a/source/tests/pt_expt/test_multitask.py b/source/tests/pt_expt/test_multitask.py
index 478f0cf6aa..2842be6dde 100644
--- a/source/tests/pt_expt/test_multitask.py
+++ b/source/tests/pt_expt/test_multitask.py
@@ -395,6 +395,9 @@ def test_multitask_train(self) -> None:
                         f"case_embd should NOT be shared: {state_key}",
                     )
 
+        if self.descriptor.get("type") == "dpa3":
+            self._assert_frozen_heads()
+
     def test_multitask_finetune(self) -> None:
         """Train, then finetune with 4 branches from pretrained 2-branch model.
 
@@ -704,6 +707,200 @@ def _get_case_embd(mk):
         # Run 1 step to verify no crash
         trainer_ft.run()
 
+    def test_multitask_finetune_no_change_model_params(self) -> None:
+        """Test change_model_params=False preserves user config in multi-task finetune.
+
+        Contrasts with change_model_params=True which overwrites descriptor/fitting
+        from pretrained (preserving trainable flags).
+        """
+        from deepmd.pt_expt.utils.finetune import (
+            get_finetune_rules,
+        )
+
+        # Phase 1: train 2-branch multi-task model (2 steps)
+        config_pretrain = _make_multitask_config(
+            self.descriptor, share_fitting=self.share_fitting, numb_steps=2
+        )
+        config_pretrain["training"]["save_freq"] = 2
+        config_pretrain["model"], shared_links_pre = preprocess_shared_params(
+            config_pretrain["model"]
+        )
+        config_pretrain = update_deepmd_input(config_pretrain, warning=False)
+        config_pretrain = normalize(config_pretrain, multi_task=True)
+        trainer = get_trainer(config_pretrain, shared_links=shared_links_pre)
+        trainer.run()
+
+        ckpt_path = os.path.join(os.getcwd(), "model.ckpt.pt")
+        self.assertTrue(os.path.exists(ckpt_path), "Pretrained checkpoint not created")
+
+        # Phase 2: build finetune config with marker in descriptor
+        ft_config = _make_multitask_config(
+            self.descriptor, share_fitting=self.share_fitting
+        )
+        # model_1: no finetune_head → resume (model_1 in pretrained_keys)
+        # model_2: finetune_head="model_2" → finetune
+        ft_config["model"]["model_dict"]["model_2"]["finetune_head"] = "model_2"
+        ft_config["training"]["numb_steps"] = 1
+        ft_config["training"]["save_freq"] = 1
+
+        # Add markers to descriptor in each branch (before preprocess_shared_params
+        # resolves shared_dict references)
+        ft_config["model"]["shared_dict"]["my_descriptor"]["_test_marker"] = True
+
+        # Phase 3: test change_model_params=False
+        ft_config_false = deepcopy(ft_config)
+        ft_config_false["model"], _ = preprocess_shared_params(ft_config_false["model"])
+        model_config_false, finetune_links_false = get_finetune_rules(
+            ckpt_path, deepcopy(ft_config_false["model"]), change_model_params=False
+        )
+
+        # User config preserved: marker still present
+        self.assertTrue(
+            model_config_false["model_dict"]["model_1"]["descriptor"].get(
+                "_test_marker", False
+            ),
+            "model_1 descriptor should preserve _test_marker with change_model_params=False",
+        )
+        self.assertTrue(
+            model_config_false["model_dict"]["model_2"]["descriptor"].get(
+                "_test_marker", False
+            ),
+            "model_2 descriptor should preserve _test_marker with change_model_params=False",
+        )
+        # FinetuneRuleItem has correct type_map
+        for mk in ("model_1", "model_2"):
+            self.assertEqual(
+                finetune_links_false[mk].get_finetune_tmap(),
+                ["O", "H"],
+                f"{mk} finetune tmap should be ['O','H']",
+            )
+        # model_1 is resuming, model_2 is not
+        self.assertTrue(
+            finetune_links_false["model_1"].resuming,
+            "model_1 should be resuming (no finetune_head, name in pretrained_keys)",
+        )
+        self.assertFalse(
+            finetune_links_false["model_2"].resuming,
+            "model_2 should not be resuming (has finetune_head)",
+        )
+
+        # Phase 4: test change_model_params=True (contrast)
+        ft_config_true = deepcopy(ft_config)
+        # Also set trainable=False to verify it's preserved
+        ft_config_true["model"]["shared_dict"]["my_descriptor"]["trainable"] = False
+        ft_config_true["model"], _ = preprocess_shared_params(ft_config_true["model"])
+        model_config_true, _finetune_links_true = get_finetune_rules(
+            ckpt_path, deepcopy(ft_config_true["model"]), change_model_params=True
+        )
+
+        # Marker overwritten from pretrained
+        self.assertFalse(
+            model_config_true["model_dict"]["model_1"]["descriptor"].get(
+                "_test_marker", False
+            ),
+            "model_1 descriptor should NOT have _test_marker with change_model_params=True",
+        )
+        self.assertFalse(
+            model_config_true["model_dict"]["model_2"]["descriptor"].get(
+                "_test_marker", False
+            ),
+            "model_2 descriptor should NOT have _test_marker with change_model_params=True",
+        )
+        # trainable=False should be preserved
+        self.assertFalse(
+            model_config_true["model_dict"]["model_1"]["descriptor"].get(
+                "trainable", True
+            ),
+            "model_1 descriptor trainable should be preserved as False",
+        )
+        self.assertFalse(
+            model_config_true["model_dict"]["model_2"]["descriptor"].get(
+                "trainable", True
+            ),
+            "model_2 descriptor trainable should be preserved as False",
+        )
+
+        # Phase 5: build trainer with change_model_params=False → run 1 step
+        ft_config_run = deepcopy(ft_config)
+        ft_config_run["model"], shared_links_ft = preprocess_shared_params(
+            ft_config_run["model"]
+        )
+        ft_config_run["model"], finetune_links_run = get_finetune_rules(
+            ckpt_path, ft_config_run["model"], change_model_params=False
+        )
+        ft_config_run = update_deepmd_input(ft_config_run, warning=False)
+        ft_config_run = normalize(ft_config_run, multi_task=True)
+        trainer_ft = get_trainer(
+            deepcopy(ft_config_run),
+            finetune_model=ckpt_path,
+            shared_links=shared_links_ft,
+            finetune_links=finetune_links_run,
+        )
+        trainer_ft.run()
+
+    def test_multitask_restart(self) -> None:
+        """Train, then restart from checkpoint and verify."""
+        # Phase 1: train
+        config1 = deepcopy(self.config)
+        config1["training"]["numb_steps"] = 2
+        config1["training"]["save_freq"] = 2
+        trainer1 = get_trainer(config1, shared_links=self.shared_links)
+        trainer1.run()
+
+        ckpt_path = "model.ckpt.pt"
+        self.assertTrue(os.path.exists(ckpt_path), "Checkpoint not created")
+
+        # Phase 2: restart to step 4
+        config2 = deepcopy(self.config)
+        config2["training"]["numb_steps"] = 4
+        config2["training"]["save_freq"] = 4
+        trainer2 = get_trainer(
+            config2,
+            restart_model=ckpt_path,
+            shared_links=self.shared_links,
+        )
+        self.assertEqual(trainer2.start_step, 2)
+        trainer2.run()
+
+    def _assert_frozen_heads(self) -> None:
+        """Export a trained head and reject missing or invalid head selections."""
+        from deepmd.pt_expt.entrypoints.main import (
+            freeze,
+        )
+        from deepmd.pt_expt.model import (
+            BaseModel,
+        )
+        from deepmd.pt_expt.utils.serialization import (
+            serialize_from_file,
+        )
+
+        checkpoint = "model.ckpt.pt"
+        output = "frozen_model_1.pte"
+        freeze(model=checkpoint, output=output, head="model_1")
+        self.assertTrue(os.path.exists(output), "Frozen model not created")
+        data = serialize_from_file(output)
+        self.assertIn("model", data)
+        self.assertIsInstance(BaseModel.deserialize(data["model"]), torch.nn.Module)
+        with self.assertRaises(ValueError, msg="Should require --head"):
+            freeze(model=checkpoint, output="frozen.pte", head=None)
+        with self.assertRaises(ValueError, msg="Should reject invalid head"):
+            freeze(model=checkpoint, output="frozen.pte", head="nonexistent")
+
+    def tearDown(self) -> None:
+        for f in os.listdir("."):
+            if f.startswith("model") and f.endswith(".pt"):
+                os.remove(f)
+            if f == "lcurve.out":
+                os.remove(f)
+            if f.endswith(".pte"):
+                os.remove(f)
+        if os.path.isdir("stat_files"):
+            shutil.rmtree("stat_files")
+
+
+class SingleTaskFinetuneTest:
+    """Fine-tune independent task heads from a single-head checkpoint."""
+
     def test_multitask_finetune_from_single_task(self) -> None:
         """Finetune multi-task model from a single-task pretrained .pt checkpoint.
 
@@ -713,11 +910,6 @@ def test_multitask_finetune_from_single_task(self) -> None:
         model_1: finetune_head="Default" → copies from single-task pretrained
         model_2: no finetune_head, not in pretrained_keys=["Default"] → RANDOM fitting
         """
-        if self.share_fitting:
-            # Single-task pretrained has no dim_case_embd; incompatible with
-            # shared fitting multi-task target.
-            return
-
         from deepmd.pt_expt.utils.finetune import (
             get_finetune_rules,
         )
@@ -858,136 +1050,9 @@ def test_multitask_finetune_from_single_task(self) -> None:
         # Phase 5: run 1 step to verify no crash
         trainer_ft.run()
 
-    def test_multitask_finetune_no_change_model_params(self) -> None:
-        """Test change_model_params=False preserves user config in multi-task finetune.
 
-        Contrasts with change_model_params=True which overwrites descriptor/fitting
-        from pretrained (preserving trainable flags).
-        """
-        from deepmd.pt_expt.utils.finetune import (
-            get_finetune_rules,
-        )
-
-        # Phase 1: train 2-branch multi-task model (2 steps)
-        config_pretrain = _make_multitask_config(
-            self.descriptor, share_fitting=self.share_fitting, numb_steps=2
-        )
-        config_pretrain["training"]["save_freq"] = 2
-        config_pretrain["model"], shared_links_pre = preprocess_shared_params(
-            config_pretrain["model"]
-        )
-        config_pretrain = update_deepmd_input(config_pretrain, warning=False)
-        config_pretrain = normalize(config_pretrain, multi_task=True)
-        trainer = get_trainer(config_pretrain, shared_links=shared_links_pre)
-        trainer.run()
-
-        ckpt_path = os.path.join(os.getcwd(), "model.ckpt.pt")
-        self.assertTrue(os.path.exists(ckpt_path), "Pretrained checkpoint not created")
-
-        # Phase 2: build finetune config with marker in descriptor
-        ft_config = _make_multitask_config(
-            self.descriptor, share_fitting=self.share_fitting
-        )
-        # model_1: no finetune_head → resume (model_1 in pretrained_keys)
-        # model_2: finetune_head="model_2" → finetune
-        ft_config["model"]["model_dict"]["model_2"]["finetune_head"] = "model_2"
-        ft_config["training"]["numb_steps"] = 1
-        ft_config["training"]["save_freq"] = 1
-
-        # Add markers to descriptor in each branch (before preprocess_shared_params
-        # resolves shared_dict references)
-        ft_config["model"]["shared_dict"]["my_descriptor"]["_test_marker"] = True
-
-        # Phase 3: test change_model_params=False
-        ft_config_false = deepcopy(ft_config)
-        ft_config_false["model"], _ = preprocess_shared_params(ft_config_false["model"])
-        model_config_false, finetune_links_false = get_finetune_rules(
-            ckpt_path, deepcopy(ft_config_false["model"]), change_model_params=False
-        )
-
-        # User config preserved: marker still present
-        self.assertTrue(
-            model_config_false["model_dict"]["model_1"]["descriptor"].get(
-                "_test_marker", False
-            ),
-            "model_1 descriptor should preserve _test_marker with change_model_params=False",
-        )
-        self.assertTrue(
-            model_config_false["model_dict"]["model_2"]["descriptor"].get(
-                "_test_marker", False
-            ),
-            "model_2 descriptor should preserve _test_marker with change_model_params=False",
-        )
-        # FinetuneRuleItem has correct type_map
-        for mk in ("model_1", "model_2"):
-            self.assertEqual(
-                finetune_links_false[mk].get_finetune_tmap(),
-                ["O", "H"],
-                f"{mk} finetune tmap should be ['O','H']",
-            )
-        # model_1 is resuming, model_2 is not
-        self.assertTrue(
-            finetune_links_false["model_1"].resuming,
-            "model_1 should be resuming (no finetune_head, name in pretrained_keys)",
-        )
-        self.assertFalse(
-            finetune_links_false["model_2"].resuming,
-            "model_2 should not be resuming (has finetune_head)",
-        )
-
-        # Phase 4: test change_model_params=True (contrast)
-        ft_config_true = deepcopy(ft_config)
-        # Also set trainable=False to verify it's preserved
-        ft_config_true["model"]["shared_dict"]["my_descriptor"]["trainable"] = False
-        ft_config_true["model"], _ = preprocess_shared_params(ft_config_true["model"])
-        model_config_true, _finetune_links_true = get_finetune_rules(
-            ckpt_path, deepcopy(ft_config_true["model"]), change_model_params=True
-        )
-
-        # Marker overwritten from pretrained
-        self.assertFalse(
-            model_config_true["model_dict"]["model_1"]["descriptor"].get(
-                "_test_marker", False
-            ),
-            "model_1 descriptor should NOT have _test_marker with change_model_params=True",
-        )
-        self.assertFalse(
-            model_config_true["model_dict"]["model_2"]["descriptor"].get(
-                "_test_marker", False
-            ),
-            "model_2 descriptor should NOT have _test_marker with change_model_params=True",
-        )
-        # trainable=False should be preserved
-        self.assertFalse(
-            model_config_true["model_dict"]["model_1"]["descriptor"].get(
-                "trainable", True
-            ),
-            "model_1 descriptor trainable should be preserved as False",
-        )
-        self.assertFalse(
-            model_config_true["model_dict"]["model_2"]["descriptor"].get(
-                "trainable", True
-            ),
-            "model_2 descriptor trainable should be preserved as False",
-        )
-
-        # Phase 5: build trainer with change_model_params=False → run 1 step
-        ft_config_run = deepcopy(ft_config)
-        ft_config_run["model"], shared_links_ft = preprocess_shared_params(
-            ft_config_run["model"]
-        )
-        ft_config_run["model"], finetune_links_run = get_finetune_rules(
-            ckpt_path, ft_config_run["model"], change_model_params=False
-        )
-        ft_config_run = update_deepmd_input(ft_config_run, warning=False)
-        ft_config_run = normalize(ft_config_run, multi_task=True)
-        trainer_ft = get_trainer(
-            deepcopy(ft_config_run),
-            finetune_model=ckpt_path,
-            shared_links=shared_links_ft,
-            finetune_links=finetune_links_run,
-        )
-        trainer_ft.run()
+class MixedTypesStatTest:
+    """Preserve and extend descriptor statistics for mixed-type systems."""
 
     def test_change_type_map_stat(self) -> None:
         """Validate change_type_map preserves existing types' stats.
@@ -999,9 +1064,6 @@ def test_change_type_map_stat(self) -> None:
            new type gets data-computed values (davg=0, dstd=0.1 for zero
            observations via StatItem defaults).
         """
-        if not self.mixed_types:
-            return
-
         old_tmap = ["O", "H"]
         new_tmap = ["O", "H", "B"]
 
@@ -1141,117 +1203,6 @@ def _is_std_like(k: str) -> bool:
                         )
                     break
 
-    def test_multitask_restart(self) -> None:
-        """Train, then restart from checkpoint and verify."""
-        # Phase 1: train
-        config1 = deepcopy(self.config)
-        config1["training"]["numb_steps"] = 2
-        config1["training"]["save_freq"] = 2
-        trainer1 = get_trainer(config1, shared_links=self.shared_links)
-        trainer1.run()
-
-        ckpt_path = "model.ckpt.pt"
-        self.assertTrue(os.path.exists(ckpt_path), "Checkpoint not created")
-
-        # Phase 2: restart to step 4
-        config2 = deepcopy(self.config)
-        config2["training"]["numb_steps"] = 4
-        config2["training"]["save_freq"] = 4
-        trainer2 = get_trainer(
-            config2,
-            restart_model=ckpt_path,
-            shared_links=self.shared_links,
-        )
-        self.assertEqual(trainer2.start_step, 2)
-        trainer2.run()
-
-    def test_multitask_freeze(self) -> None:
-        """Train, then freeze with --head and verify.
-
-        Only runs for dpa3 descriptor to avoid redundant slow freeze tests.
-        """
-        if self.descriptor.get("type") != "dpa3":
-            return
-
-        from deepmd.pt_expt.entrypoints.main import (
-            freeze,
-        )
-
-        # Train
-        config = deepcopy(self.config)
-        trainer = get_trainer(config, shared_links=self.shared_links)
-        trainer.run()
-
-        # Freeze head model_1
-        ckpt_path = "model.ckpt.pt"
-        output_path = "frozen_model_1.pte"
-        freeze(model=ckpt_path, output=output_path, head="model_1")
-        self.assertTrue(os.path.exists(output_path), "Frozen model not created")
-
-        # Verify frozen model loads
-        from deepmd.pt_expt.model import (
-            BaseModel,
-        )
-        from deepmd.pt_expt.utils.serialization import (
-            serialize_from_file,
-        )
-
-        data = serialize_from_file(output_path)
-        self.assertIn("model", data)
-        frozen_model = BaseModel.deserialize(data["model"])
-        self.assertIsInstance(frozen_model, torch.nn.Module)
-
-    def test_multitask_freeze_no_head_raises(self) -> None:
-        """Freezing multi-task model without --head raises ValueError.
-
-        Only runs for dpa3 descriptor to avoid redundant slow freeze tests.
-        """
-        if self.descriptor.get("type") != "dpa3":
-            return
-
-        from deepmd.pt_expt.entrypoints.main import (
-            freeze,
-        )
-
-        config = deepcopy(self.config)
-        trainer = get_trainer(config, shared_links=self.shared_links)
-        trainer.run()
-
-        ckpt_path = "model.ckpt.pt"
-        with self.assertRaises(ValueError, msg="Should require --head"):
-            freeze(model=ckpt_path, output="frozen.pte", head=None)
-
-    def test_multitask_freeze_invalid_head_raises(self) -> None:
-        """Freezing multi-task model with invalid --head raises ValueError.
-
-        Only runs for dpa3 descriptor to avoid redundant slow freeze tests.
-        """
-        if self.descriptor.get("type") != "dpa3":
-            return
-
-        from deepmd.pt_expt.entrypoints.main import (
-            freeze,
-        )
-
-        config = deepcopy(self.config)
-        trainer = get_trainer(config, shared_links=self.shared_links)
-        trainer.run()
-
-        ckpt_path = "model.ckpt.pt"
-        with self.assertRaises(ValueError, msg="Should reject invalid head"):
-            freeze(model=ckpt_path, output="frozen.pte", head="nonexistent")
-
-    def tearDown(self) -> None:
-        for f in os.listdir("."):
-            if f.startswith("model") and f.endswith(".pt"):
-                os.remove(f)
-            if f == "lcurve.out":
-                os.remove(f)
-            if f.endswith(".pte"):
-                os.remove(f)
-        if os.path.isdir("stat_files"):
-            shutil.rmtree("stat_files")
-
 
 class TestMultiTaskEpochSchedule(unittest.TestCase):
     """Test the run length and task weights derived from num_epoch_dict."""
@@ -1289,7 +1240,11 @@ def test_steps_and_prob_follow_epoch_targets(self) -> None:
         np.testing.assert_allclose(trainer.model_prob, [0.25, 0.75])
 
 
-class TestMultiTaskSeA(unittest.TestCase, MultiTaskTrainTest):
+class TestMultiTaskSeA(
+    unittest.TestCase,
+    MultiTaskTrainTest,
+    SingleTaskFinetuneTest,
+):
     """Multi-task training with se_e2_a descriptor."""
 
     @classmethod
@@ -1315,7 +1270,10 @@ def tearDown(self) -> None:
         shutil.rmtree(self.tmpdir, ignore_errors=True)
 
 
-class TestMultiTaskSeAShareFit(unittest.TestCase, MultiTaskTrainTest):
+class TestMultiTaskSeAShareFit(
+    unittest.TestCase,
+    MultiTaskTrainTest,
+):
     """Multi-task training with se_e2_a descriptor and shared fitting_net."""
 
     @classmethod
@@ -1341,7 +1299,12 @@ def tearDown(self) -> None:
         shutil.rmtree(self.tmpdir, ignore_errors=True)
 
 
-class TestMultiTaskDPA1(unittest.TestCase, MultiTaskTrainTest):
+class TestMultiTaskDPA1(
+    unittest.TestCase,
+    MultiTaskTrainTest,
+    SingleTaskFinetuneTest,
+    MixedTypesStatTest,
+):
     """Multi-task training with DPA1 (se_atten) descriptor."""
 
     @classmethod
@@ -1367,7 +1330,11 @@ def tearDown(self) -> None:
         shutil.rmtree(self.tmpdir, ignore_errors=True)
 
 
-class TestMultiTaskDPA1ShareFit(unittest.TestCase, MultiTaskTrainTest):
+class TestMultiTaskDPA1ShareFit(
+    unittest.TestCase,
+    MultiTaskTrainTest,
+    MixedTypesStatTest,
+):
     """Multi-task training with DPA1 descriptor and shared fitting_net."""
 
     @classmethod
@@ -1393,7 +1360,12 @@ def tearDown(self) -> None:
         shutil.rmtree(self.tmpdir, ignore_errors=True)
 
 
-class TestMultiTaskDPA2(unittest.TestCase, MultiTaskTrainTest):
+class TestMultiTaskDPA2(
+    unittest.TestCase,
+    MultiTaskTrainTest,
+    SingleTaskFinetuneTest,
+    MixedTypesStatTest,
+):
     """Multi-task training with DPA2 descriptor."""
 
     @classmethod
@@ -1419,7 +1391,11 @@ def tearDown(self) -> None:
         shutil.rmtree(self.tmpdir, ignore_errors=True)
 
 
-class TestMultiTaskDPA2ShareFit(unittest.TestCase, MultiTaskTrainTest):
+class TestMultiTaskDPA2ShareFit(
+    unittest.TestCase,
+    MultiTaskTrainTest,
+    MixedTypesStatTest,
+):
     """Multi-task training with DPA2 descriptor and shared fitting_net."""
 
     @classmethod
@@ -1445,7 +1421,12 @@ def tearDown(self) -> None:
         shutil.rmtree(self.tmpdir, ignore_errors=True)
 
 
-class TestMultiTaskDPA3(unittest.TestCase, MultiTaskTrainTest):
+class TestMultiTaskDPA3(
+    unittest.TestCase,
+    MultiTaskTrainTest,
+    SingleTaskFinetuneTest,
+    MixedTypesStatTest,
+):
     """Multi-task training with DPA3 descriptor."""
 
     @classmethod
@@ -1471,7 +1452,11 @@ def tearDown(self) -> None:
         shutil.rmtree(self.tmpdir, ignore_errors=True)
 
 
-class TestMultiTaskDPA3ShareFit(unittest.TestCase, MultiTaskTrainTest):
+class TestMultiTaskDPA3ShareFit(
+    unittest.TestCase,
+    MultiTaskTrainTest,
+    MixedTypesStatTest,
+):
     """Multi-task training with DPA3 descriptor and shared fitting_net."""
 
     @classmethod
@@ -1890,20 +1875,6 @@ def _make_gradient_test_st_config(
     }
 
 
-def _deterministic_task_choice(task_sequence: list[int]):
-    """Return a patched dp_random.choice that forces task selection order."""
-    original = dp_random.choice
-    it = iter(task_sequence)
-
-    def patched(a, size=None, replace=True, p=None):
-        # Task selection: array with >=2 elements and probability vector
-        if hasattr(a, "__len__") and len(a) >= 2 and p is not None:
-            return next(it)
-        return original(a, size=size, replace=replace, p=p)
-
-    return patched
-
-
 def _make_recording_step(
     trainer,
     modules_to_record: dict,
@@ -2003,9 +1974,12 @@ def test_gradient_accumulation(self) -> None:
             {"desc": mt_desc, "fit_1": mt_fit_1, "fit_2": mt_fit_2},
             mt_grads,
         )
-        with patch(
-            "deepmd.utils.random.choice",
-            _deterministic_task_choice([0, 1]),
+        with patch.object(
+            mt_trainer,
+            "select_task",
+            side_effect=[
+                mt_trainer.training_tasks[key] for key in ("model_1", "model_2")
+            ],
         ):
             mt_trainer.run()  # 2 steps: task_0 then task_1
 
@@ -2150,8 +2124,8 @@ class TestCompileCaseEmbdVaryingNframes(unittest.TestCase):
     The test uses two systems with different atom counts and per-system
     ``batch_size=[2, 3]`` so every branch's compiled graph sees both
     nframes values.  ``dim_case_embd=2`` is deliberately chosen to also
-    collide numerically with the nframes=2 runtime case.  ``dp_random.choice``
-    is mocked so both tasks and both systems are sampled.
+    collide numerically with the nframes=2 runtime case. Tasks are selected
+    explicitly, and ``dp_random.choice`` controls data-system selection.
     """
 
     @classmethod
diff --git a/source/tests/pt_expt/test_train_gradient.py b/source/tests/pt_expt/test_train_gradient.py
index f20912f57b..69c3425d92 100644
--- a/source/tests/pt_expt/test_train_gradient.py
+++ b/source/tests/pt_expt/test_train_gradient.py
@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import unittest
+from unittest.mock import (
+    patch,
+)
 
 import torch
 
@@ -87,6 +90,32 @@ def test_finite_norms_do_not_raise(self) -> None:
     def test_no_update_is_noop(self) -> None:
         NonFiniteGradGuard().raise_if_nonfinite(self._named(1.0))
 
+    def test_remote_failure_is_agreed_before_diagnostics(self) -> None:
+        # Even a rank with no local update must join the agreement. In FSDP,
+        # accessing the gradients for diagnostics may itself be collective.
+        guard = NonFiniteGradGuard()
+        events = []
+
+        def agree(gathered, local_flag):
+            self.assertFalse(local_flag)
+            events.append("agree")
+            gathered[:] = [False, True]
+
+        def named_parameters():
+            events.append("diagnostics")
+            return self._named(1.0)()
+
+        with (
+            patch("deepmd.pt_expt.train.gradient.dist.get_world_size", return_value=2),
+            patch(
+                "deepmd.pt_expt.train.gradient.dist.all_gather_object",
+                side_effect=agree,
+            ),
+            self.assertRaisesRegex(RuntimeError, "Non-finite gradient norm"),
+        ):
+            guard.raise_if_nonfinite(named_parameters, synchronize=True)
+        self.assertEqual(events, ["agree", "diagnostics"])
+
     def test_reports_offending_parameter(self) -> None:
         guard = NonFiniteGradGuard()
         guard.update(torch.tensor(float("nan"), device="cpu"))
diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py
index fa4199195b..fd87f8a0ea 100644
--- a/source/tests/pt_expt/test_training.py
+++ b/source/tests/pt_expt/test_training.py
@@ -765,69 +765,6 @@ def test_missing_attr_raises(self) -> None:
             _ = cm.nonexistent_attribute_xyz
 
 
-@REQUIRES_SUPPORTED_COMPILE
-class TestCompiledDynamicShapes(unittest.TestCase):
-    """Test that _CompiledModel handles varying nall via dynamic shapes."""
-
-    @classmethod
-    def setUpClass(cls) -> None:
-        data_dir = os.path.join(EXAMPLE_DIR, "data")
-        if not os.path.isdir(data_dir):
-            raise unittest.SkipTest(f"Example data not found: {data_dir}")
-        cls.data_dir = data_dir
-
-    def test_compiled_handles_varying_nall(self) -> None:
-        """Run several training steps, assert finite loss each step.
-
-        With ``tracing_mode="symbolic"`` + ``dynamic=True``, nall is a
-        symbolic dim so nall growth across batches is handled without
-        any recompile or padding.
-        """
-        from deepmd.pt_expt.train.training import (
-            _CompiledModel,
-        )
-
-        config = _make_config(self.data_dir, numb_steps=5)
-        config["training"]["enable_compile"] = True
-        config = update_deepmd_input(config, warning=False)
-        config = normalize(config)
-
-        tmpdir = tempfile.mkdtemp(prefix="pt_expt_dynamic_")
-        try:
-            old_cwd = os.getcwd()
-            os.chdir(tmpdir)
-            try:
-                trainer = get_trainer(config)
-
-                # The wrapper.model should be a _CompiledModel
-                compiled_model = trainer.wrapper.model["Default"]
-                self.assertIsInstance(compiled_model, _CompiledModel)
-                # Lazy compile: no graph exists before any forward.
-                self.assertEqual(compiled_model._compiled_lower_by_mode, {})
-
-                trainer.wrapper.train()
-                for step in range(3):
-                    trainer.optimizer.zero_grad(set_to_none=True)
-                    inp, lab = trainer.get_data(is_train=True)
-                    lr = trainer.scheduler.get_last_lr()[0]
-                    _, loss, _ = trainer.wrapper(**inp, cur_lr=lr, label=lab)
-                    loss.backward()
-                    trainer.optimizer.step()
-
-                    # After the first forward the training graph must exist,
-                    # keyed by the mode it was traced in.
-                    if step == 0:
-                        self.assertIn(True, compiled_model._compiled_lower_by_mode)
-
-                    # Loss should be a finite scalar at every step
-                    self.assertFalse(torch.isnan(loss))
-                    self.assertFalse(torch.isinf(loss))
-            finally:
-                os.chdir(old_cwd)
-        finally:
-            shutil.rmtree(tmpdir, ignore_errors=True)
-
-
 @REQUIRES_SUPPORTED_COMPILE
 class TestCompiledConsistency(unittest.TestCase):
     """Verify compiled model produces the same energy/force/virial as uncompiled."""
@@ -1112,7 +1049,7 @@ def test_replicated_numpy_filter_reads_a_complete_local_pass(
     trainer.multi_task = False
     trainer.world_size = world_size
     trainer.training_data_by_task = {"Default": data}
-    trainer.min_pair_dist_by_task = {"Default": 0.5}
+    trainer.pair_filter_enabled_by_task = {"Default": True}
     with (
         patch("deepmd.utils.data.dp_random.shuffle"),
         patch.object(data, "get_batch", wraps=data.get_batch) as get_batch,
@@ -2018,6 +1955,7 @@ def _check_varying_natoms(self, descriptor: dict | None = None) -> None:
                 trainer_c = get_trainer(self._make_varying_config(True, descriptor))
                 compiled_model = trainer_c.wrapper.model["Default"]
                 self.assertIsInstance(compiled_model, _CompiledModel)
+                self.assertEqual(compiled_model._compiled_lower_by_mode, {})
 
                 # Sync weights so predictions can be compared exactly
                 compiled_model.original_model.load_state_dict(
@@ -2044,6 +1982,8 @@ def _check_varying_natoms(self, descriptor: dict | None = None) -> None:
                         out_c, loss_c, _ = trainer_c.wrapper(
                             **inp, cur_lr=lr, label=lab
                         )
+                        self.assertIn(True, compiled_model._compiled_lower_by_mode)
+                        self.assertTrue(torch.isfinite(loss_c))
 
                         ctx = f"step={step} "
                         _assert_compile_predictions_match(self, out_c, out_uc, ctx=ctx)
diff --git a/source/tests/pt_expt/utils/test_canonical_graph.py b/source/tests/pt_expt/utils/test_canonical_graph.py
index 3128db789a..b59a71fee3 100644
--- a/source/tests/pt_expt/utils/test_canonical_graph.py
+++ b/source/tests/pt_expt/utils/test_canonical_graph.py
@@ -6,9 +6,12 @@
     NeighborGraph,
 )
 from deepmd.pt_expt.utils.canonical_graph import (
-    UINT32_MAX,
+    INT32_MAX,
     CanonicalGraph,
     canonical_graph_from_neighbor_graph,
+    canonical_graph_from_pairs,
+    canonical_index_bytes,
+    canonical_index_words,
     validate_canonical_graph_shapes,
 )
 
@@ -45,18 +48,21 @@ def test_storage_guards_remain_outside_csr(physical_edge_count: int) -> None:
     assert compact.edge_vec.shape == (2, 3)
     assert compact.source_order.shape == (2,)
     assert compact.source.dtype == torch.uint32
-    assert compact.source_order.dtype == torch.uint32
+    assert compact.source_order.dtype == torch.int64
     assert int(compact.destination_row_ptr[-1]) == physical_edge_count
     assert int(compact.source_row_ptr[-1]) == physical_edge_count
     validate_canonical_graph_shapes(compact, 1)
 
 
-def test_storage_exceeding_uint32_range_is_rejected() -> None:
-    storage_count = UINT32_MAX + 1
+@pytest.mark.parametrize("index_dtype", [torch.uint32, torch.int64])
+def test_edge_storage_range_is_independent_of_node_width(
+    index_dtype: torch.dtype,
+) -> None:
+    storage_count = 1 << 32
     graph = CanonicalGraph(
         n_node=torch.empty(1, dtype=torch.int64, device="meta"),
         n_local=torch.empty(1, dtype=torch.int64, device="meta"),
-        source=torch.empty(storage_count, dtype=torch.uint32, device="meta"),
+        source=torch.empty(storage_count, dtype=index_dtype, device="meta"),
         edge_vec=torch.empty(
             storage_count,
             3,
@@ -67,9 +73,96 @@ def test_storage_exceeding_uint32_range_is_rejected() -> None:
         source_row_ptr=torch.empty(2, dtype=torch.int64, device="meta"),
         source_order=torch.empty(
             storage_count,
-            dtype=torch.uint32,
+            dtype=torch.int64,
             device="meta",
         ),
     )
-    with pytest.raises(ValueError, match="exceeds the uint32 range"):
-        validate_canonical_graph_shapes(graph, 1)
+    validate_canonical_graph_shapes(graph, 1)
+
+
+@pytest.mark.parametrize(
+    "node_count, expected_dtype",
+    [(INT32_MAX + 1, torch.uint32), (INT32_MAX + 2, torch.int64)],
+)
+def test_pairs_choose_width_from_node_extent(
+    node_count: int, expected_dtype: torch.dtype
+) -> None:
+    graph = canonical_graph_from_pairs(
+        torch.empty(2, dtype=torch.int64, device="meta"),
+        torch.empty(2, dtype=torch.int64, device="meta"),
+        torch.empty(2, 3, dtype=torch.float32, device="meta"),
+        torch.empty(1, dtype=torch.int64, device="meta"),
+        node_count,
+    )
+    assert graph.source.dtype == expected_dtype
+    assert graph.source_order.dtype == torch.int64
+    validate_canonical_graph_shapes(graph, node_count)
+
+
+@pytest.mark.parametrize("dtype", [torch.uint32, torch.int32, torch.int64])
+def test_index_byte_carrier_preserves_words_without_copy(dtype: torch.dtype) -> None:
+    values = [0, 7, (1 << 40) + 123] if dtype == torch.int64 else [0, 7, INT32_MAX]
+    index = torch.tensor(values, dtype=dtype)
+    storage = canonical_index_bytes(index)
+    restored = canonical_index_words(storage)
+    assert storage.shape == (3, index.element_size())
+    assert storage.data_ptr() == restored.data_ptr() == index.data_ptr()
+    torch.testing.assert_close(
+        restored.to(torch.int64), index.to(torch.int64), atol=0, rtol=0
+    )
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("edge_count", [0, 1, 5])
+def test_pairs_preserve_geometry_and_both_csr_views(
+    device: str, edge_count: int
+) -> None:
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    center = torch.tensor([3, 1, 0, 3, 1], device=device)[:edge_count]
+    source = torch.tensor([2, 0, 1, 2, 1], device=device)[:edge_count]
+    vectors = torch.arange(15, dtype=torch.float64, device=device).reshape(5, 3)
+    graph = canonical_graph_from_pairs(
+        center,
+        source,
+        vectors[:edge_count],
+        torch.tensor([2, 99, 2, 99], device=device)[::2],
+        4,
+    )
+    validate_canonical_graph_shapes(graph, 4)
+    torch.testing.assert_close(
+        graph.destination_row_ptr.diff(), torch.bincount(center, minlength=4)
+    )
+    torch.testing.assert_close(
+        graph.source_row_ptr.diff(), torch.bincount(source, minlength=4)
+    )
+    destination = torch.repeat_interleave(
+        torch.arange(4, device=device), graph.destination_row_ptr.diff()
+    )
+    payload = (
+        torch.cat(
+            (
+                destination[:, None],
+                graph.source[:edge_count].to(torch.int64)[:, None],
+                graph.edge_vec[:edge_count],
+            ),
+            dim=1,
+        )
+        .cpu()
+        .tolist()
+    )
+    expected = (
+        torch.cat((center[:, None], source[:, None], vectors[:edge_count]), dim=1)
+        .cpu()
+        .tolist()
+    )
+    assert sorted(payload) == sorted(expected)
+    by_source = graph.source.to(torch.int64).index_select(
+        0, graph.source_order[:edge_count].to(torch.int64)
+    )
+    torch.testing.assert_close(
+        by_source,
+        torch.repeat_interleave(
+            torch.arange(4, device=device), graph.source_row_ptr.diff()
+        ),
+    )
diff --git a/source/tests/pt_expt/utils/test_common.py b/source/tests/pt_expt/utils/test_common.py
index 57ee153e54..179426c929 100644
--- a/source/tests/pt_expt/utils/test_common.py
+++ b/source/tests/pt_expt/utils/test_common.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
 import numpy as np
+import pytest
 import torch
 
 from deepmd.dpmodel.common import (
@@ -127,3 +128,61 @@ class MockModule(MockNativeOP):
     assert hasattr(module, "forward")
     # Should not have forward_lower method since there's no call_lower method
     assert not hasattr(module, "forward_lower")
+
+
+def test_dpmodel_setattr_keeps_module_device(monkeypatch) -> None:
+    """NumPy buffer assignment after Module.to keeps the module device.
+
+    ``env.DEVICE`` is patched to ``meta`` so a regression that always places
+    new buffers on the process default would land off-CPU and fail the assert.
+    """
+    from deepmd.dpmodel.common import (
+        NativeOP,
+    )
+    from deepmd.pt_expt.common import (
+        torch_module,
+    )
+
+    class HolderDP(NativeOP):
+        def __init__(self) -> None:
+            self.stat = np.zeros(3, dtype=np.float64)
+
+        def call(self, x: np.ndarray) -> np.ndarray:
+            return x
+
+    @torch_module
+    class Holder(HolderDP):
+        pass
+
+    module = Holder().to("cpu")
+    assert module.stat.device.type == "cpu"
+    monkeypatch.setattr(env, "DEVICE", torch.device("meta"))
+    module.stat = np.ones(3, dtype=np.float64)
+    assert module.stat.device.type == "cpu"
+    np.testing.assert_array_equal(module.stat.detach().cpu().numpy(), np.ones(3))
+
+
+@pytest.mark.parametrize("resident", ["parameter", "buffer"])  # device source
+def test_dpmodel_setattr_uses_other_module_tensors(monkeypatch, resident) -> None:
+    """A new NumPy buffer follows an existing parameter or other buffer."""
+
+    @torch_module
+    class Holder(NativeOP):
+        def call(self, *args, **kwargs) -> None:
+            pass
+
+    monkeypatch.setattr(env, "DEVICE", torch.device("cpu"))
+    module = Holder()
+    tensor = torch.zeros(3, device="cpu")
+    if resident == "parameter":
+        module.weight = torch.nn.Parameter(tensor)
+    else:
+        module.register_buffer("other", tensor)
+        assert not list(module.parameters())
+    module.to("meta")
+
+    assert "stat" not in module._buffers
+    module.stat = np.ones(3, dtype=np.float64)
+    assert module.stat.device.type == "meta"
+    assert module.stat.dtype == torch.float64
+    assert module.stat.shape == (3,)
diff --git a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py
index 74d7b17246..4d7c5f3751 100644
--- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py
+++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py
@@ -21,9 +21,16 @@
 import pytest
 import torch
 
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+    ModelOutputDef,
+)
 from deepmd.entrypoints.convert_backend import (
     convert_backend,
 )
+from deepmd.pt_expt.infer.deep_eval import (
+    DeepEval,
+)
 from deepmd.pt_expt.model.graph_lower import (
     graph_edge_dtype,
 )
@@ -493,9 +500,11 @@ def test_dense_pt2_has_lower_input_kind_nlist(dpa1_dpmodel_data) -> None:
 def test_neighbor_graph_method_rejected_on_nlist_artifact(dpa1_dpmodel_data) -> None:
     """A non-default ``neighbor_graph_method`` on a NLIST-form artifact raises.
 
-    The knob is consumed only by graph-form ``.pt2`` eval; silently ignoring
-    it on nlist-form artifacts misled users into thinking they selected an
-    O(N) builder. The nlist-path knob is ``nlist_backend``.
+    The knob is consumed only by the artifacts whose inputs the evaluator
+    builds from a neighbor graph -- the graph, compact canonical and
+    edge-vector lowers; silently ignoring it on nlist-form artifacts misled
+    users into thinking they selected an O(N) builder. The nlist-path knob is
+    ``nlist_backend``.
     """
     from deepmd.infer import (
         DeepPot,
@@ -508,12 +517,40 @@ def test_neighbor_graph_method_rejected_on_nlist_artifact(dpa1_dpmodel_data) ->
             copy.deepcopy(dpa1_dpmodel_data),
             do_atomic_virial=True,
         )
-        with pytest.raises(ValueError, match="graph-form"):
+        with pytest.raises(ValueError, match="use nlist_backend"):
             DeepPot(p, neighbor_graph_method="vesin")
         # The default remains valid for nlist-form artifacts.
         DeepPot(p)
 
 
+@pytest.mark.parametrize(
+    ("method", "backend", "message"),
+    [
+        ("dense", "auto", "Edge-vector artifacts support neighbor_graph_method"),
+        ("vesin", "native", "Select only one"),
+    ],
+)
+def test_edge_vec_neighbor_options_rejected(
+    monkeypatch: pytest.MonkeyPatch, method: str, backend: str, message: str
+) -> None:
+    """The constructor validates edge options independently of compiled loading."""
+
+    def load_metadata(evaluator: DeepEval, model_file: str) -> None:
+        evaluator.metadata = {"lower_input_kind": "edge_vec"}
+
+    monkeypatch.setattr(DeepEval, "_load_pt2", load_metadata)
+    with (
+        pytest.warns(DeprecationWarning, match="deprecated edge_vec lower schema"),
+        pytest.raises(ValueError, match=message),
+    ):
+        DeepEval(
+            "edge.pt2",
+            ModelOutputDef(FittingOutputDef([])),
+            neighbor_graph_method=method,
+            nlist_backend=backend,
+        )
+
+
 class _FakeDesc:
     def __init__(self, n_attn: int) -> None:
         self._n = n_attn
diff --git a/source/tests/pt_expt/utils/test_graph_trace_samples.py b/source/tests/pt_expt/utils/test_graph_trace_samples.py
new file mode 100644
index 0000000000..cc9f7df9aa
--- /dev/null
+++ b/source/tests/pt_expt/utils/test_graph_trace_samples.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Graph trace coverage must not depend on dense neighbor capacities."""
+
+from unittest.mock import (
+    patch,
+)
+
+import pytest
+import torch
+
+from deepmd.pt_expt.utils.serialization import (
+    build_synthetic_graph_inputs,
+    count_synthetic_graph_edges,
+)
+
+
+class _GraphSampleModel:
+    def __init__(self, ntypes=118):
+        self.ntypes = ntypes
+
+    def get_rcut(self):
+        return 4.0
+
+    def get_type_map(self):
+        return [f"T{i}" for i in range(self.ntypes)]
+
+    def get_sel(self):
+        raise AssertionError("Graph samples must not query dense capacities")
+
+    def get_dim_fparam(self):
+        return 0
+
+    def get_dim_aparam(self):
+        return 0
+
+    def get_dim_chg_spin(self):
+        return 0
+
+
+@pytest.mark.parametrize("nloc", [4, 7])
+@pytest.mark.parametrize("ntypes", [1, 118])
+@pytest.mark.parametrize("canonicalize", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_graph_sample_covers_sparse_topology_and_high_types(
+    nloc, ntypes, canonicalize, dtype
+):
+    model = _GraphSampleModel(ntypes)
+    kwargs = {"nframes": 2, "nloc": nloc, "dtype": dtype, "device": torch.device("cpu")}
+    with patch(
+        "deepmd.pt_expt.utils.serialization.build_neighbor_list",
+        side_effect=AssertionError("Graph samples must not build dense lists"),
+    ):
+        real = count_synthetic_graph_edges(model, **kwargs)
+        sample = build_synthetic_graph_inputs(
+            model, real + 5, canonicalize=canonicalize, **kwargs
+        )
+        repeated = build_synthetic_graph_inputs(
+            model, real + 5, canonicalize=canonicalize, **kwargs
+        )
+    for actual, expected in zip(sample, repeated, strict=True):
+        if actual is not None:
+            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    atype, n_node, n_local, edge_index, edge_vec, edge_mask = sample[:6]
+    assert atype.shape == (2 * nloc,)
+    assert atype[0] == atype[nloc] == ntypes - 1
+    assert (atype < ntypes - 1).any() == (ntypes > 1)
+    assert n_node.tolist() == [nloc, nloc]
+    assert n_local.tolist() == [nloc - 1, nloc - 1]
+    assert edge_mask.sum() == real
+    assert (~edge_mask).sum() >= 5
+    real_edges = edge_index[:, edge_mask]
+    # An owned center in each frame has no incoming or outgoing real edge.
+    assert not torch.isin(real_edges, torch.tensor([1, nloc + 1])).any()
+    assert 1 < n_local.min()
+    # The highest-type node has two guaranteed neighbors and occurs at
+    # both ends of directed edges.
+    assert (real_edges[0] == 0).sum() >= 2
+    assert (real_edges[1] == 0).sum() >= 2
+    assert torch.isfinite(edge_vec).all()
+    assert all(value is not None for value in sample[6:10]) == canonicalize
+
+
+def test_graph_export_high_type_sparse_runtime():
+    from deepmd.pt_expt.model import (
+        get_model,
+    )
+    from deepmd.pt_expt.utils import (
+        env,
+    )
+    from deepmd.pt_expt.utils.serialization import (
+        _trace_and_export,
+    )
+
+    config = {
+        "type_map": [f"T{i}" for i in range(118)],
+        "descriptor": {
+            "type": "se_atten",
+            "sel": 416,
+            "rcut": 4.0,
+            "rcut_smth": 0.5,
+            "neuron": [3, 6],
+            "axis_neuron": 2,
+            "attn_layer": 0,
+            "seed": 42,
+        },
+        "fitting_net": {"neuron": [8, 8], "seed": 42},
+    }
+    model = get_model(config).to(env.DEVICE).eval()
+    exported, metadata, _, _ = _trace_and_export(
+        {"model": model.serialize()}, model_json_override=None, lower_kind="graph"
+    )
+    # Dense compatibility metadata remains available to existing C++ readers.
+    assert metadata["sel"] == model.get_sel()
+    assert metadata["nnei"] == sum(model.get_sel())
+    loaded = exported.module()
+    for nframes, nloc in ((1, 1), (1, 4), (3, 8)):
+        inputs = list(
+            build_synthetic_graph_inputs(
+                model,
+                None,
+                nframes=nframes,
+                nloc=nloc,
+                dtype=torch.float64,
+                device=env.DEVICE,
+            )
+        )
+        inputs[2] = inputs[1].clone()  # Runtime may own every atom.
+        actual = loaded(*inputs)
+        reference = model.forward_common_lower_graph(*inputs[:10])
+        torch.testing.assert_close(
+            actual["energy"], reference["energy_redu"], rtol=1e-10, atol=1e-10
+        )
+        for key, ref_key in (
+            ("force", "energy_derv_r"),
+            ("virial", "energy_derv_c_redu"),
+        ):
+            torch.testing.assert_close(
+                actual[key],
+                reference[ref_key].reshape(actual[key].shape),
+                rtol=1e-10,
+                atol=1e-10,
+            )
diff --git a/source/tests/pt_expt/utils/test_network.py b/source/tests/pt_expt/utils/test_network.py
index 54f12554c1..08bec4683b 100644
--- a/source/tests/pt_expt/utils/test_network.py
+++ b/source/tests/pt_expt/utils/test_network.py
@@ -98,9 +98,12 @@ def test_serialization_round_trip_pt_expt(self) -> None:
         net2 = EmbeddingNet.deserialize(serialized)
 
         # Verify layers are still pt_expt NativeLayer modules
-        self.assertIsInstance(net2.layers, torch.nn.ModuleList)
-        for layer in net2.layers:
-            self.assertIsInstance(layer, NativeLayer)
+        for network in (net, net2):
+            self.assertIsInstance(network.layers, torch.nn.ModuleList)
+            for layer in network.layers:
+                self.assertIsInstance(layer, NativeLayer)
+                self.assertIsInstance(layer, torch.nn.Module)
+                self.assertTrue(hasattr(layer, "_parameters"))
 
         out2 = net2(x)
         np.testing.assert_allclose(
@@ -108,40 +111,6 @@ def test_serialization_round_trip_pt_expt(self) -> None:
             out2.detach().cpu().numpy(),
         )
 
-    def test_deserialize_preserves_layer_type(self) -> None:
-        """Test that deserialize uses type(obj.layers[0]) to preserve subclass layers.
-
-        This is the key fix: dpmodel's deserialize no longer hardcodes
-        super(EmbeddingNet, obj).__init__(layers), which would overwrite
-        pt_expt's converted layers. Instead it uses type(obj.layers[0])
-        to respect the subclass's layer type.
-        """
-        # Create pt_expt EmbeddingNet
-        net = EmbeddingNet(
-            in_dim=self.in_dim,
-            neuron=self.neuron,
-            activation_function=self.activation,
-            resnet_dt=self.resnet_dt,
-            precision=self.precision,
-            seed=GLOBAL_SEED,
-        )
-
-        # Verify layers are pt_expt NativeLayer (torch modules)
-        for layer in net.layers:
-            self.assertIsInstance(layer, torch.nn.Module)
-            self.assertTrue(hasattr(layer, "_parameters"))
-
-        # Deserialize
-        serialized = net.serialize()
-        net2 = EmbeddingNet.deserialize(serialized)
-
-        # Verify deserialized layers are STILL pt_expt NativeLayer, not dpmodel
-        for layer in net2.layers:
-            self.assertIsInstance(layer, torch.nn.Module)
-            self.assertTrue(hasattr(layer, "_parameters"))
-            # This would fail if deserialize used hardcoded dpmodel layers
-            self.assertIsInstance(layer, NativeLayer)
-
     def test_cross_backend_consistency(self) -> None:
         """Test numerical consistency between dpmodel and pt_expt EmbeddingNet."""
         # Create both with same seed
diff --git a/source/tests/pt_expt/utils/test_nv_graph_builder.py b/source/tests/pt_expt/utils/test_nv_graph_builder.py
index b8e5b7dc41..9288de7e0f 100644
--- a/source/tests/pt_expt/utils/test_nv_graph_builder.py
+++ b/source/tests/pt_expt/utils/test_nv_graph_builder.py
@@ -14,6 +14,10 @@
     build_neighbor_graph,
 )
 
+from ...common.dpmodel.test_neighbor_graph import (
+    assert_csr_consistent,
+)
+
 nv_builder = pytest.importorskip("deepmd.pt_expt.utils.nv_graph_builder")
 from deepmd.pt.utils.nv_nlist import (
     is_nv_available,
@@ -113,7 +117,9 @@ def test_nv_excludes_virtual_atoms_like_dense():
     assert np.all(at[ei[0]] >= 0) and np.all(at[ei[1]] >= 0)
 
 
-def test_nv_pair_exclusion_matches_dense():
+@pytest.mark.parametrize("csr_mode", ["none", "views", "canonical"])
+@pytest.mark.parametrize("exclude_pairs", [False, True])
+def test_nv_pair_exclusion_matches_dense(csr_mode: str, exclude_pairs: bool) -> None:
     dev = torch.device("cuda")
     coord = torch.tensor(
         [[[0.0, 0.0, 0.0], [0.8, 0.0, 0.0], [0.0, 1.1, 0.0], [1.2, 1.0, 0.0]]],
@@ -122,7 +128,7 @@ def test_nv_pair_exclusion_matches_dense():
     )
     atype = torch.tensor([[0, 1, 0, 1]], dtype=torch.int64, device=dev)
     box = (torch.eye(3, dtype=torch.float64, device=dev) * 3.0).reshape(1, 3, 3)
-    pair_excl = PairExcludeMask(2, [(0, 1), (1, 0)])
+    pair_excl = PairExcludeMask(2, [(0, 1), (1, 0)]) if exclude_pairs else None
     expected = build_neighbor_graph(
         coord,
         atype,
@@ -136,5 +142,9 @@ def test_nv_pair_exclusion_matches_dense():
         box,
         2.0,
         pair_excl=pair_excl,
+        with_csr=csr_mode != "none",
+        canonicalize=csr_mode == "canonical",
     )
     assert _edge_multiset(actual) == _edge_multiset(expected)
+    if csr_mode != "none":
+        assert_csr_consistent(actual)
diff --git a/source/tests/pt_expt/utils/test_serialization_kernel_levels.py b/source/tests/pt_expt/utils/test_serialization_kernel_levels.py
index 53651103ca..0816908353 100644
--- a/source/tests/pt_expt/utils/test_serialization_kernel_levels.py
+++ b/source/tests/pt_expt/utils/test_serialization_kernel_levels.py
@@ -1,11 +1,19 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Kernel-level selection for pt_expt serialization."""
 
+import json
 import os
 import sys
+import zipfile
+from pathlib import (
+    Path,
+)
 from types import (
     ModuleType,
 )
+from unittest import (
+    mock,
+)
 
 import pytest
 import torch
@@ -261,3 +269,41 @@ def test_target_policy_only_suppresses_incompatible_dpa4_accelerators(
         assert tuple(os.environ[name] for name in _INFER_LEVELS) == expected
 
     assert all(os.environ[name] == "1" for name in _INFER_LEVELS)
+
+
+@pytest.mark.parametrize("single_process", [False, True])
+def test_communication_artifact_is_optional(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, single_process: bool
+) -> None:
+    """The default includes communication; a single-process export omits it."""
+    data = {"model": {"type": "dpa4"}}
+
+    def trace(*args: object, **kwargs: object) -> tuple:
+        return (
+            None,
+            {"has_comm_artifact": True, "has_chg_spin_ebd": False},
+            data,
+            ["energy"],
+        )
+
+    def compile_package(exported: object, package_path: str, **kwargs: object) -> None:
+        with zipfile.ZipFile(package_path, "w") as archive:
+            archive.writestr("model/data.pkl", b"")
+
+    monkeypatch.setattr(serialization, "_trace_and_export", trace)
+    output = tmp_path / "model.pt2"
+    kwargs = {"with_comm": False} if single_process else {}
+    with mock.patch(
+        "torch._inductor.aoti_compile_and_package", side_effect=compile_package
+    ) as compile_mock:
+        serialization.deserialize_to_file(
+            str(output), data, lower_kind="graph", **kwargs
+        )
+
+    assert compile_mock.call_count == (1 if single_process else 2)
+    with zipfile.ZipFile(output) as archive:
+        metadata = json.loads(archive.read("model/extra/metadata.json"))
+        assert metadata["has_comm_artifact"] is (not single_process)
+        assert ("model/extra/forward_lower_with_comm.pt2" in archive.namelist()) is (
+            not single_process
+        )
diff --git a/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py b/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py
new file mode 100644
index 0000000000..7d3e372236
--- /dev/null
+++ b/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py
@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""The eager CPU vacuum fold must not execute deployment-only operators."""
+
+import copy
+import os
+
+import pytest
+import torch
+from torch.fx.experimental.proxy_tensor import (
+    make_fx,
+)
+
+from deepmd.pt_expt.kernels import (
+    autotune,
+)
+from deepmd.pt_expt.kernels.cuda.dpa4 import (
+    edge_radial,
+)
+from deepmd.pt_expt.model.get_model import (
+    get_model,
+)
+from deepmd.pt_expt.model.model import (
+    BaseModel,
+)
+from deepmd.pt_expt.utils import (
+    env,
+    serialization,
+)
+
+from ..model.test_dpa4_vacuum_ref import (
+    CONFIG,
+    PRESET,
+)
+
+LEVELS = (
+    "DP_CUDA_INFER",
+    "DP_TRITON_INFER",
+    "DP_CUTILE_INFER",
+    "DP_CUTE_INFER",
+    "DP_CUDA_TRAIN",
+    "DP_TRITON_TRAIN",
+    "DP_AMP_INFER",
+)
+
+
+@torch.library.custom_op(
+    "deepmd_test::vacuum_cuda_radial", mutates_args=(), device_types="cuda"
+)
+def cuda_radial(
+    edge_len: torch.Tensor, keep: torch.Tensor, freqs: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """A CUDA-only operator, including on test hosts without a CUDA build."""
+    raise AssertionError("the eager fold must never call the accelerator")
+
+
+@cuda_radial.register_fake
+def cuda_radial_fake(edge_len, keep, freqs):
+    return edge_len.new_empty(edge_len.shape), edge_len.new_empty(
+        (edge_len.shape[0], freqs.numel())
+    )
+
+
+def reference_model(monkeypatch, numb_fparam=0, numb_aparam=0, vacuum_ref=True):
+    monkeypatch.setattr(env, "DEVICE", torch.device("cpu"))
+    for name in LEVELS:
+        monkeypatch.setenv(name, "0")
+    config = copy.deepcopy(CONFIG)
+    config["descriptor"].update(precision="float32")
+    config["fitting_net"].update(
+        precision="float32",
+        vacuum_ref=vacuum_ref,
+        numb_fparam=numb_fparam,
+        numb_aparam=numb_aparam,
+    )
+    config["preset_out_bias"] = PRESET
+    return get_model(config).eval()
+
+
+@pytest.mark.parametrize(
+    ("numb_fparam", "numb_aparam"), [(0, 0), (1, 0), (0, 1)]
+)  # frame/atomic parameter dimensions
+def test_cuda_export_folds_on_cpu_before_restoring_deployment_dispatch(
+    monkeypatch, numb_fparam, numb_aparam
+) -> None:
+    """Construction-time CUDA bindings survive an eager, no-grad reference fold."""
+    reference = reference_model(monkeypatch, numb_fparam, numb_aparam)
+    data = {"model": reference.serialize()}
+    reference.fold_vacuum_reference()
+    expected = reference.atomic_model.fitting_net
+
+    # Exercise the production constructor gate without requiring CUDA hardware
+    # or its extension. This op has a fake implementation but no CPU kernel,
+    # exactly the distinction between tracing and the eager vacuum evaluation.
+    monkeypatch.setattr(
+        edge_radial,
+        "make_cuda_edge_radial",
+        lambda envelope, basis: (
+            lambda distance, keep: cuda_radial(distance, keep, basis.adam_freqs)
+        ),
+    )
+    monkeypatch.setenv("DP_CUDA_INFER", "1")
+    model = BaseModel.deserialize(data["model"]).eval()
+    am = model.atomic_model
+    descriptor = am.descriptor
+    assert descriptor._cuda_radial_fn is not None
+    with pytest.raises(NotImplementedError, match="CPU"):
+        model.fold_vacuum_reference()
+
+    # A CUDA target is restored before the exporter tunes and traces. Setting
+    # it only after construction also makes this test work on CPU-only hosts.
+    monkeypatch.setattr(env, "DEVICE", torch.device("cuda"))
+    monkeypatch.setattr(BaseModel, "deserialize", lambda data: model)
+    events = []
+    move = model.to
+    fold = model.fold_vacuum_reference
+
+    def move_to_cpu(device):
+        events.append("move")
+        assert device == "cpu"
+        return move(device)
+
+    def fold_on_cpu():
+        events.append("fold")
+        assert env.DEVICE.type == "cpu"
+        assert not torch.is_grad_enabled()
+        assert am.descriptor is not descriptor
+        assert not am.descriptor.training
+        assert am.descriptor._cuda_radial_fn is None
+        assert all(t.device.type == "cpu" for t in model.parameters())
+        assert all(t.device.type == "cpu" for t in model.buffers())
+        fold()
+
+    class StopBeforeTrace(Exception):
+        pass
+
+    def check_before_trace(model, target, lower_kind):
+        events.append("trace")
+        assert target.type == env.DEVICE.type == "cuda"
+        assert torch.is_grad_enabled(), "force tracing still needs autograd"
+        assert am.descriptor is descriptor
+        assert os.environ["DP_CUDA_INFER"] == "1"
+        raise StopBeforeTrace
+
+    monkeypatch.setattr(model, "to", move_to_cpu)
+    monkeypatch.setattr(model, "fold_vacuum_reference", fold_on_cpu)
+    monkeypatch.setattr(autotune, "run_autotune", check_before_trace)
+    with pytest.raises(StopBeforeTrace):
+        serialization._trace_and_export_impl(data, target_device=env.DEVICE)
+    assert events == ["move", "fold", "trace"]
+
+    fitting = am.fitting_net
+    assert fitting.vacuum_ref == expected.vacuum_ref
+    torch.testing.assert_close(fitting.bias_atom_e, expected.bias_atom_e)
+    torch.testing.assert_close(fitting.vacuum_table, expected.vacuum_table)
+    if numb_fparam or numb_aparam:
+        assert fitting.vacuum_table.device.type == "cpu"
+        assert "vacuum_table" not in fitting.state_dict()
+    assert all(t.device.type == "cpu" for t in model.buffers())
+
+    # Restoring the descriptor must preserve optimized fake tracing, rather
+    # than silently exporting the reference implementation for the real atoms.
+    with torch.no_grad():
+        traced = make_fx(
+            lambda: am.vacuum_descriptor(),
+            tracing_mode="fake",
+            _allow_non_fake_inputs=True,
+        )()
+    assert any(node.target == cuda_radial._opoverload for node in traced.graph.nodes)
+
+
+@pytest.mark.parametrize("failure", ["construction", "fold"])  # failing stage
+def test_reference_context_restores_state_after_failure(monkeypatch, failure) -> None:
+    model = reference_model(monkeypatch)
+    am = model.atomic_model
+    descriptor = am.descriptor
+    monkeypatch.setattr(env, "DEVICE", torch.device("meta"))
+    for i, name in enumerate(LEVELS):
+        if i % 2:
+            monkeypatch.delenv(name)
+        else:
+            monkeypatch.setenv(name, "1")
+    saved = {name: os.environ.get(name) for name in LEVELS}
+
+    def fail(data):
+        raise RuntimeError("reference construction failed")
+
+    if failure == "construction":
+        monkeypatch.setattr(type(descriptor), "deserialize", fail)
+    with pytest.raises(RuntimeError, match="failed"):
+        with serialization._cpu_reference_descriptors(model):
+            assert am.descriptor is not descriptor
+            assert env.DEVICE.type == "cpu"
+            assert all(os.environ[name] == "0" for name in LEVELS)
+            raise RuntimeError("reference fold failed")
+    assert am.descriptor is descriptor
+    assert env.DEVICE.type == "meta"
+    assert {name: os.environ.get(name) for name in LEVELS} == saved
+
+
+@pytest.mark.parametrize("folded_table", [False, True])  # disabled vs already folded
+def test_unneeded_reference_does_not_rebuild_descriptor(
+    monkeypatch, folded_table
+) -> None:
+    model = reference_model(monkeypatch, numb_fparam=1, vacuum_ref=folded_table)
+    if folded_table:
+        model.fold_vacuum_reference()
+    descriptor = model.atomic_model.descriptor
+
+    def fail(data):
+        raise AssertionError("an unreferenced descriptor should not be rebuilt")
+
+    monkeypatch.setattr(type(descriptor), "deserialize", fail)
+    monkeypatch.setattr(env, "DEVICE", torch.device("meta"))
+    with serialization._cpu_reference_descriptors(model):
+        assert model.atomic_model.descriptor is descriptor
+        assert env.DEVICE.type == "meta"
+
+
+def test_reference_descriptor_preserves_disabled_graph_route(monkeypatch) -> None:
+    """Virtual-spin backbones disable graph lowering outside serialization."""
+    monkeypatch.setattr(env, "DEVICE", torch.device("cpu"))
+    for name in LEVELS:
+        monkeypatch.setenv(name, "0")
+    model = get_model(
+        {
+            "type_map": ["O", "H"],
+            "descriptor": {
+                "type": "dpa1",
+                "rcut": 4.0,
+                "rcut_smth": 0.5,
+                "sel": [8],
+                "attn_layer": 0,
+                "axis_neuron": 2,
+                "neuron": [4, 8],
+                "seed": 1,
+            },
+            "fitting_net": {
+                "type": "ener",
+                "neuron": [8],
+                "vacuum_ref": True,
+                "seed": 1,
+            },
+            "preset_out_bias": PRESET,
+        }
+    ).eval()
+    am = model.atomic_model
+    descriptor = am.descriptor
+    descriptor.disable_graph_lower()
+    with torch.no_grad():
+        descriptor.se_atten.mean.fill_(0.3)
+        expected = am.vacuum_descriptor()
+        with serialization._cpu_reference_descriptors(model):
+            assert not am.descriptor.uses_graph_lower()
+            torch.testing.assert_close(am.vacuum_descriptor(), expected)
+    assert am.descriptor is descriptor
diff --git a/source/tests/pt_expt/utils/test_tabulate_ops_lazy.py b/source/tests/pt_expt/utils/test_tabulate_ops_lazy.py
index 841ada65c3..9b52b9b103 100644
--- a/source/tests/pt_expt/utils/test_tabulate_ops_lazy.py
+++ b/source/tests/pt_expt/utils/test_tabulate_ops_lazy.py
@@ -23,6 +23,9 @@
 import textwrap
 
 import torch
+from torch._subclasses.fake_tensor import (
+    FakeTensorMode,
+)
 
 from deepmd.pt_expt.utils import (
     tabulate_ops,
@@ -105,3 +108,21 @@ def _fallback(*args, **kwargs):
             ns.__dict__.pop(op_name, None)
         tabulate_ops._registered.clear()
         tabulate_ops._registered.update(saved_registered)
+
+
+def test_composite_operator_uses_its_decomposition_for_fake_tensors() -> None:
+    """Composite operators retain PyTorch's fake-tensor decomposition."""
+    with torch.library._scoped_library("deepmd_test_tabulate", "DEF") as library:
+        library.define("identity(Tensor x) -> Tensor")
+        library.impl("identity", lambda x: x + 1, "CompositeImplicitAutograd")
+
+        def fake(value: torch.Tensor) -> torch.Tensor:
+            raise AssertionError("Composite operators must use their decomposition")
+
+        tabulate_ops._try_register_fake("deepmd_test_tabulate::identity", fake)
+        with FakeTensorMode():
+            value = torch.ones((3, 4), device="cpu")
+            result = torch.ops.deepmd_test_tabulate.identity(value)
+        assert result.shape == value.shape
+        assert result.dtype == value.dtype
+        assert result.device == value.device
diff --git a/source/tests/pt_expt/utils/test_vesin_graph_builder.py b/source/tests/pt_expt/utils/test_vesin_graph_builder.py
index c216676994..57d1acf3a3 100644
--- a/source/tests/pt_expt/utils/test_vesin_graph_builder.py
+++ b/source/tests/pt_expt/utils/test_vesin_graph_builder.py
@@ -1,4 +1,11 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from dataclasses import (
+    replace,
+)
+from unittest import (
+    mock,
+)
+
 import numpy as np
 import pytest
 import torch
@@ -10,11 +17,22 @@
     apply_pair_exclusion,
     build_neighbor_graph,
 )
+from deepmd.pt_expt.utils.canonical_graph import (
+    canonical_graph_from_neighbor_graph,
+    validate_canonical_graph_shapes,
+)
+
+from ...common.dpmodel.test_neighbor_graph import (
+    assert_csr_consistent,
+)
 
 vesin_builder = pytest.importorskip("deepmd.pt_expt.utils.vesin_graph_builder")
 from deepmd.pt_expt.utils.vesin_neighbor_list import (
     is_vesin_torch_available,
 )
+from deepmd.pt_expt.utils.vesin_search import (
+    VesinNeighborSearch,
+)
 
 pytestmark = pytest.mark.skipif(
     not is_vesin_torch_available(), reason="vesin[torch] not installed"
@@ -113,9 +131,9 @@ def test_vesin_excludes_virtual_atoms_like_dense():
 
 def _valid_edge_set(ng):
     """Return the set of (src, dst, rounded edge_vec) for all real edges."""
-    ei = np.asarray(ng.edge_index)
-    ev = np.asarray(ng.edge_vec)
-    em = np.asarray(ng.edge_mask)
+    ei = ng.edge_index.detach().cpu().numpy()
+    ev = ng.edge_vec.detach().cpu().numpy()
+    em = ng.edge_mask.detach().cpu().numpy()
     return {
         (int(ei[0, k]), int(ei[1, k]), tuple(np.round(ev[k], 6)))
         for k in range(ei.shape[1])
@@ -137,7 +155,8 @@ def test_vesin_pair_excl_none_identity(periodic):
 
 
 @pytest.mark.parametrize("periodic", [False, True])  # non-PBC and PBC
-def test_vesin_pair_excl_oracle_set_equality(periodic):
+@pytest.mark.parametrize("csr_mode", ["none", "views", "canonical"])
+def test_vesin_pair_excl_oracle_set_equality(periodic: bool, csr_mode: str) -> None:
     """Vesin builder(pair_excl=X) == dense ref + apply_pair_exclusion(X)."""
     coord, atype, box = _system(periodic)
     coord = coord.reshape(1, 4, 3)
@@ -150,9 +169,17 @@ def test_vesin_pair_excl_oracle_set_equality(periodic):
     ng_ref = apply_pair_exclusion(ng_dense, atype_flat, pe)
     # vesin builder with fused post-process
     ng_vesin = vesin_builder.build_neighbor_graph_vesin(
-        coord, atype, box_3d, rcut, pair_excl=pe
+        coord,
+        atype,
+        box_3d,
+        rcut,
+        pair_excl=pe,
+        with_csr=csr_mode != "none",
+        canonicalize=csr_mode == "canonical",
     )
     assert _valid_edge_set(ng_ref) == _valid_edge_set(ng_vesin)
+    if csr_mode != "none":
+        assert_csr_consistent(ng_vesin)
     # exclusion actually removed edges
     ng_plain = vesin_builder.build_neighbor_graph_vesin(coord, atype, box_3d, rcut)
     assert int(np.asarray(ng_vesin.edge_mask).sum()) < int(
@@ -172,3 +199,200 @@ def test_vesin_nlist_edges_pair_excl_raises():
     nl = VesinNeighborList()
     with pytest.raises(NotImplementedError, match="return_mode='edges'"):
         nl.build(coord, atype, None, 2.0, [4], return_mode="edges", pair_excl=pe)
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_cached_graph_tracks_geometry_and_atom_changes(device: str) -> None:
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    coord, atype, box = _system(True)
+    coord = coord.reshape(1, 4, 3).to(device)
+    atype = atype.to(device)
+    box = box.reshape(1, 3, 3).to(device)
+    search = VesinNeighborSearch(skin=1.0, sorted=True)
+    for step in range(9):
+        points = coord.clone()
+        types = atype.clone()
+        cell = box.clone()
+        if step == 1:
+            points[0, 1, 0] += 0.2
+        elif step == 2:
+            points[0, 1, 0] += 0.8
+        elif step == 3:
+            cell *= 1.15
+        elif step == 4:
+            cell = None
+        elif step == 5:
+            types[0, 1] = -1
+        elif step == 6:
+            points = torch.cat((points, points[:, :1] + 0.4), dim=1)
+            types = torch.cat((types, types[:, :1]), dim=1)
+        elif step == 7:
+            cell[0, 0, 1] += 0.3
+        elif step == 8:
+            points[0, 1] += cell[0, 0]
+        actual = vesin_builder.build_neighbor_graph_vesin(
+            points, types, cell, 2.0, search=search, canonicalize=True
+        )
+        expected = build_neighbor_graph(points, types, cell, 2.0, canonicalize=True)
+        assert _valid_edge_set(actual) == _valid_edge_set(expected)
+        torch.testing.assert_close(actual.n_node, expected.n_node)
+        assert actual.destination_sorted
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_cached_search_preserves_earlier_autograd_inputs(device: str) -> None:
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    coord, atype, box = _system(True)
+    coord = coord.reshape(1, 4, 3).to(device).requires_grad_()
+    atype = atype.to(device)
+    box = box.reshape(1, 3, 3).to(device)
+    search = VesinNeighborSearch(skin=1.0, sorted=True)
+    graph = vesin_builder.build_neighbor_graph_vesin(
+        coord, atype, box, 2.0, search=search
+    )
+    moved = coord.detach().clone()
+    moved[0, 1, 0] += 0.3
+    vesin_builder.build_neighbor_graph_vesin(moved, atype, box, 2.0, search=search)
+    actual = torch.autograd.grad(graph.edge_vec.square().sum(), coord)[0]
+    reference = build_neighbor_graph(coord, atype, box, 2.0)
+    expected = torch.autograd.grad(reference.edge_vec.square().sum(), coord)[0]
+    torch.testing.assert_close(actual, expected)
+
+
+def test_cached_search_keeps_frames_independent() -> None:
+    coord, atype, box = _system(True)
+    coord = coord.unsqueeze(0).repeat(2, 1, 1)
+    atype = atype.repeat(2, 1)
+    box = box.unsqueeze(0).repeat(2, 1, 1)
+    box[1] *= 1.5
+    search = VesinNeighborSearch(skin=1.0, sorted=True)
+    for _ in range(3):
+        actual = vesin_builder.build_neighbor_graph_vesin(
+            coord, atype, box, 2.0, search=search, canonicalize=True
+        )
+        expected = build_neighbor_graph(coord, atype, box, 2.0, canonicalize=True)
+        assert _valid_edge_set(actual) == _valid_edge_set(expected)
+        coord[1, 1, 0] += 0.2
+
+
+@pytest.mark.parametrize("return_vectors", [False, True])
+def test_cached_search_rebuilds_for_cell_contraction(return_vectors: bool) -> None:
+    """Cell contraction includes images absent from the previous candidates."""
+    from vesin_torch import (
+        NeighborList,
+    )
+
+    points = torch.zeros((1, 3), dtype=torch.float64)
+    cell = torch.diag(torch.tensor([2.35, 10.0, 10.0], dtype=torch.float64))
+    other_cell = cell.clone()
+    search = VesinNeighborSearch(skin=1.0)
+    with mock.patch("vesin_torch.NeighborList", wraps=NeighborList) as constructor:
+        for frame in (0, 1):
+            ii, _, _ = search.compute(points, cell, 6.0, frame=frame)
+            assert ii.numel() == 4
+        search.compute(points, cell.clone(), 6.0)
+        assert constructor.call_count == 2
+
+        # The third image moves from 7.05 to 5.85 angstrom, crossing both
+        # the old candidate radius (7.0) and the physical cutoff (6.0).
+        cell[0, 0] = 1.95
+        ii, jj, geometry = search.compute(
+            points, cell, 6.0, return_vectors=return_vectors
+        )
+        assert ii.numel() == 6
+        assert torch.count_nonzero(ii) == torch.count_nonzero(jj) == 0
+        expected = torch.tensor([-3, -2, -1, 1, 2, 3], dtype=geometry.dtype)
+        if return_vectors:
+            expected *= cell[0, 0]
+        torch.testing.assert_close(geometry[:, 0].sort().values, expected)
+        assert torch.count_nonzero(geometry[:, 1:]) == 0
+        assert constructor.call_count == 3
+
+        ii, _, _ = search.compute(points, other_cell, 6.0, frame=1)
+        assert ii.numel() == 4
+        assert constructor.call_count == 3
+        ii, _, _ = search.compute(points, None, 6.0)
+        assert ii.numel() == 0
+        ii, _, _ = search.compute(points, cell, 6.0)
+        assert ii.numel() == 6
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("periodic", [False, True])
+@pytest.mark.parametrize("nloc", [0, 1, 4])
+@pytest.mark.parametrize("nframes", [1, 2])
+def test_canonical_vesin_matches_dense(
+    device: str, periodic: bool, nloc: int, nframes: int
+) -> None:
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    coord, atype, box = _system(periodic)
+    coord = coord[:nloc].unsqueeze(0).repeat(nframes, 1, 1).to(device)
+    atype = atype[:, :nloc].repeat(nframes, 1).to(device)
+    box = box.unsqueeze(0).repeat(nframes, 1, 1).to(device) if periodic else None
+    if nloc == 4:
+        atype[0, 1] = -1
+    search = VesinNeighborSearch(skin=1.0)
+    for step in range(3):
+        points = coord.clone()
+        if nloc:
+            points[-1, 0, 0] += step * 0.3
+        actual = vesin_builder.build_canonical_graph_vesin(
+            points, box, 2.0, search, node_mask=atype >= 0 if nloc == 4 else None
+        )
+        reference = build_neighbor_graph(points, atype, box, 2.0, canonicalize=True)
+        expected = canonical_graph_from_neighbor_graph(
+            replace(reference, n_local=reference.n_node)
+        )
+        validate_canonical_graph_shapes(actual, nframes * nloc)
+        torch.testing.assert_close(actual.n_node, expected.n_node)
+        torch.testing.assert_close(
+            actual.destination_row_ptr, expected.destination_row_ptr
+        )
+        torch.testing.assert_close(actual.source_row_ptr, expected.source_row_ptr)
+        payloads = []
+        for graph in (actual, expected):
+            edge_count = int(graph.destination_row_ptr[-1])
+            destination = torch.repeat_interleave(
+                torch.arange(nframes * nloc, device=device),
+                graph.destination_row_ptr.diff(),
+            )
+            values = (
+                torch.cat(
+                    (
+                        destination[:, None],
+                        graph.source[:edge_count].to(torch.int64)[:, None],
+                        graph.edge_vec[:edge_count],
+                    ),
+                    dim=1,
+                )
+                .cpu()
+                .numpy()
+            )
+            order = np.lexsort(values[:, ::-1].T)
+            payloads.append(values[order])
+        np.testing.assert_allclose(*payloads, rtol=1e-6, atol=1e-6)
+
+
+@pytest.mark.parametrize("return_vectors", [False, True])
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_cached_search_handles_empty_nonempty_transition(
+    device: str, return_vectors: bool
+) -> None:
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    search = VesinNeighborSearch(skin=1.0)
+    for distance in (8.0, 0.8, 8.0):
+        positions = torch.tensor(
+            [[0.0, 0.0, 0.0], [distance, 0.0, 0.0]],
+            dtype=torch.float64,
+            device=device,
+        )
+        ii, jj, geometry = search.compute(
+            positions, None, 2.0, return_vectors=return_vectors
+        )
+        assert ii.numel() == (2 if distance < 2.0 else 0)
+        expected = positions[jj] - positions[ii] if return_vectors else geometry * 0
+        torch.testing.assert_close(geometry, expected)
diff --git a/source/tests/test_cc_local.py b/source/tests/test_cc_local.py
new file mode 100644
index 0000000000..f003bf55d1
--- /dev/null
+++ b/source/tests/test_cc_local.py
@@ -0,0 +1,234 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Exercise the C++ CI entry point without compiling backend libraries."""
+
+import json
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import (
+    Path,
+)
+
+import pytest
+
+pytestmark = pytest.mark.skipif(
+    sys.platform != "linux", reason="The build script requires GNU/Linux utilities"
+)
+
+SCRIPT = Path(__file__).parents[1] / "install" / "test_cc_local.sh"
+MODELS = {
+    "gen_sea.py",
+    "gen_dpa1.py",
+    "gen_dpa2.py",
+    "gen_dpa3.py",
+    "gen_fparam_aparam.py",
+    "gen_model_devi.py",
+    "gen_chg_spin.py",
+    "gen_dpa4.py",
+    "gen_dpa1_pairexcl.py",
+    "gen_spin.py",
+    "gen_spin_model_devi.py",
+    "gen_corrupt_with_comm.py",
+    "gen_dpa4_spin.py",
+    "gen_dpa4_zbl.py",
+    "gen_dpa4_spin_chgspin.py",
+    "gen_dpa4_spin_zbl.py",
+    "gen_dpa4c_spin_lmp.py",
+    "gen_dpa4c_spin.py",
+    "gen_dpa4c_canonical.py",
+    "gen_dpa4c_long_range.py",
+}
+
+
+@pytest.fixture
+def build_tree(tmp_path):
+    root = tmp_path / "repo"
+    script = root / "source/install/test_cc_local.sh"
+    script.parent.mkdir(parents=True)
+    shutil.copy2(SCRIPT, script)
+    infer = root / "source/tests/infer"
+    infer.mkdir(parents=True)
+    op = root / "source/build_tests/op/pt/libdeepmd_op_pt.so"
+    op.parent.mkdir(parents=True)
+    op.write_text("custom op from this build")
+    c_api = root / "dp_test/lib/libdeepmd_c.so"
+    c_api.parent.mkdir(parents=True)
+    c_api.touch()
+    fake_package = tmp_path / "packages/deepmd"
+    fake_package.mkdir(parents=True)
+    (fake_package / "__init__.py").touch()
+    (fake_package / "env.py").write_text(
+        "import os\nfrom pathlib import Path\n"
+        'SHARED_LIB_DIR = Path(os.environ["MOCK_SHARED_LIB_DIR"])\n'
+    )
+    bin_dir = tmp_path / "bin"
+    bin_dir.mkdir()
+    stub = f"""#!{sys.executable}
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+name = Path(sys.argv[0]).name
+args = sys.argv[1:]
+record = {{"name": name, "args": args, "ld": os.environ.get("LD_LIBRARY_PATH", "")}}
+with open(os.environ["MOCK_LOG"], "a") as stream:
+    stream.write(json.dumps(record) + "\\n")
+if os.environ.get("MOCK_FAIL") and any(os.environ["MOCK_FAIL"] == Path(a).name for a in [name, *args]):
+    sys.exit(17)
+if name == "python" and args[0] == "-c":
+    sys.exit(subprocess.call([sys.executable, "-S", *args]))
+"""
+    for name in ("cmake", "python", "ctest"):
+        executable = bin_dir / name
+        executable.write_text(stub)
+        executable.chmod(0o755)
+    env = {
+        **os.environ,
+        "PATH": str(bin_dir) + os.pathsep + os.environ["PATH"],
+        "PYTHONPATH": str(fake_package.parent),
+        "DP_VARIANT": "cuda",
+        "ENABLE_TENSORFLOW": "TRUE",
+        "ENABLE_PYTORCH": "TRUE",
+        "ENABLE_PADDLE": "FALSE",
+        "LD_LIBRARY_PATH": "/existing/lib",
+        # Do not inherit sanitizer flags from the caller.
+        "CXXFLAGS": "",
+        "MOCK_LOG": str(tmp_path / "commands.jsonl"),
+        "MOCK_SHARED_LIB_DIR": str(tmp_path / "installed_ops"),
+        "MOCK_FAIL": "",
+    }
+    return root, script, env
+
+
+def run_script(build_tree, *args, **overrides):
+    root, script, env = build_tree
+    env = {**env, **overrides}
+    result = subprocess.run(
+        ["bash", str(script), *args],
+        env=env,
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    log = Path(env["MOCK_LOG"])
+    commands = (
+        [json.loads(line) for line in log.read_text().splitlines()]
+        if log.exists()
+        else []
+    )
+    return result, commands
+
+
+def generated_models(commands):
+    return [
+        Path(c["args"][0]).name
+        for c in commands
+        if c["name"] == "python" and c["args"][0] != "-c"
+    ]
+
+
+def test_default_still_builds_all_models_and_runs_ctest(build_tree):
+    root, _, env = build_tree
+    stale = root / "source/tests/infer/stale.pt2"
+    stale.touch()
+    result, commands = run_script(build_tree)
+    assert result.returncode == 0, result.stderr
+    assert "BUILD_TESTING:BOOL=TRUE" in commands[0]["args"]
+    assert [c["name"] for c in commands[:4]] == ["cmake", "cmake", "cmake", "python"]
+    assert commands[1]["args"][:2] == ["--build", "."]
+    assert commands[2]["args"] == ["--install", "."]
+    assert commands[-1]["name"] == "ctest"
+    assert commands[-1]["args"] == ["--output-on-failure"]
+    assert set(generated_models(commands)) == MODELS
+    assert len(generated_models(commands)) == len(MODELS)
+    assert not stale.exists()
+    assert (
+        Path(env["MOCK_SHARED_LIB_DIR"]) / "libdeepmd_op_pt.so"
+    ).read_text() == "custom op from this build"
+
+
+@pytest.mark.parametrize("pytorch", ["TRUE", "FALSE"])
+def test_build_only_installs_without_generating_models_or_running_ctest(
+    build_tree, pytorch
+):
+    root, _, env = build_tree
+    stale = root / "source/tests/infer/keep.pt2"
+    stale.touch()
+    result, commands = run_script(build_tree, "--build-only", ENABLE_PYTORCH=pytorch)
+    assert result.returncode == 0, result.stderr
+    assert "BUILD_TESTING:BOOL=FALSE" in commands[0]["args"]
+    assert not generated_models(commands)
+    assert all(c["name"] != "ctest" for c in commands)
+    assert stale.exists()
+    assert (Path(env["MOCK_SHARED_LIB_DIR"]) / "libdeepmd_op_pt.so").exists() == (
+        pytorch == "TRUE"
+    )
+
+
+def test_kokkos_models_reuse_build_and_preserve_other_fixtures(build_tree):
+    root, _, _ = build_tree
+    infer = root / "source/tests/infer"
+    keep = infer / "other_model.pt2"
+    keep.touch()
+    stale = [
+        infer / f"deeppot_dpa4c_long_range_compressed.{ext}"
+        for ext in ("pt2", "pte", "expected")
+    ]
+    for path in stale:
+        path.touch()
+    result, commands = run_script(build_tree, "--models-only=kokkos")
+    assert result.returncode == 0, result.stderr
+    assert [c["name"] for c in commands] == ["python", "python"]
+    assert commands[0]["args"][0] == "-c"
+    assert generated_models(commands) == ["gen_dpa4c_long_range.py"]
+    assert (
+        Path(commands[1]["ld"].split(os.pathsep)[0]).resolve() == root / "dp_test/lib"
+    )
+    assert keep.exists()
+    assert not any(path.exists() for path in stale)
+
+
+@pytest.mark.parametrize(
+    "override", [{"DP_VARIANT": "cpu"}, {"ENABLE_PYTORCH": "FALSE"}]
+)
+def test_kokkos_models_reject_missing_cuda_or_pytorch(build_tree, override):
+    result, commands = run_script(build_tree, "--models-only=kokkos", **override)
+    assert result.returncode != 0
+    assert not commands
+
+
+@pytest.mark.parametrize(
+    "missing",
+    ["dp_test/lib/libdeepmd_c.so", "source/build_tests/op/pt/libdeepmd_op_pt.so"],
+)
+def test_kokkos_models_require_the_existing_build(build_tree, missing):
+    root, _, _ = build_tree
+    (root / missing).unlink()
+    result, commands = run_script(build_tree, "--models-only=kokkos")
+    assert result.returncode != 0
+    assert not commands
+
+
+@pytest.mark.parametrize("model", ["gen_sea.py", "gen_dpa4c_long_range.py"])
+def test_default_propagates_generator_failures(build_tree, model):
+    result, commands = run_script(build_tree, MOCK_FAIL=model)
+    assert result.returncode != 0
+    assert all(c["name"] != "ctest" for c in commands)
+
+
+def test_kokkos_models_propagate_generator_failure(build_tree):
+    result, _ = run_script(
+        build_tree, "--models-only=kokkos", MOCK_FAIL="gen_dpa4c_long_range.py"
+    )
+    assert result.returncode != 0
+
+
+@pytest.mark.parametrize(
+    "args", [("--unknown",), ("--build-only", "--models-only=kokkos")]
+)
+def test_invalid_arguments_fail_before_build(build_tree, args):
+    result, commands = run_script(build_tree, *args)
+    assert result.returncode == 2
+    assert not commands
diff --git a/source/tests/tf/test_deeppolar.py b/source/tests/tf/test_deeppolar.py
index 3d7ace564d..bd6274e57e 100644
--- a/source/tests/tf/test_deeppolar.py
+++ b/source/tests/tf/test_deeppolar.py
@@ -2,7 +2,6 @@
 import os
 import unittest
 
-import ase.neighborlist
 import numpy as np
 from packaging.version import parse as parse_version
 
@@ -1183,30 +1182,3 @@ def test_2frame_full_atm(self) -> None:
         np.testing.assert_almost_equal(
             vv.reshape([-1]), expected_gv.reshape([-1]), decimal=default_places
         )
-
-
-@unittest.skipIf(
-    parse_version(tf.__version__) < parse_version("1.15"),
-    f"The current tf version {tf.__version__} is too low to run the new testing model.",
-)
-class TestDeepPolarNewPBCNeighborList(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls) -> None:
-        convert_pbtxt_to_pb(
-            str(infer_path / os.path.join("deeppolar_new.pbtxt")),
-            "deeppolar_new.pb",
-        )
-        cls.dp = DeepPolar(
-            "deeppolar_new.pb",
-            neighbor_list=ase.neighborlist.NewPrimitiveNeighborList(
-                cutoffs=6, bothways=True
-            ),
-        )
-
-    @unittest.skip("multiple frames not supported")
-    def test_2frame_full_atm(self) -> None:
-        pass
-
-    @unittest.skip("multiple frames not supported")
-    def test_2frame_old_atm(self) -> None:
-        pass
diff --git a/source/tests/tf/test_model_compression_se_r.py b/source/tests/tf/test_model_compression_se_r.py
index c42538f05d..ad68b21aad 100644
--- a/source/tests/tf/test_model_compression_se_r.py
+++ b/source/tests/tf/test_model_compression_se_r.py
@@ -65,6 +65,28 @@ def setUpModule() -> None:
     INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models()
 
 
+def tearDownModule() -> None:
+    _file_delete(INPUT)
+    _file_delete(FROZEN_MODEL)
+    _file_delete(COMPRESSED_MODEL)
+    _file_delete("out.json")
+    _file_delete("compress.json")
+    _file_delete("checkpoint")
+    _file_delete("model.ckpt.meta")
+    _file_delete("model.ckpt.index")
+    _file_delete("model.ckpt.data-00000-of-00001")
+    _file_delete("model.ckpt-100.meta")
+    _file_delete("model.ckpt-100.index")
+    _file_delete("model.ckpt-100.data-00000-of-00001")
+    _file_delete("model-compression/checkpoint")
+    _file_delete("model-compression/model.ckpt.meta")
+    _file_delete("model-compression/model.ckpt.index")
+    _file_delete("model-compression/model.ckpt.data-00000-of-00001")
+    _file_delete("model-compression")
+    _file_delete("input_v2_compat.json")
+    _file_delete("lcurve.out")
+
+
 class TestDeepPotAPBC(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
@@ -403,149 +425,3 @@ def test_ase(self) -> None:
         nframes = 1
         np.testing.assert_almost_equal(ff0, ff1, default_places)
         np.testing.assert_almost_equal(ee0, ee1, default_places)
-
-
-class TestDeepPotAPBCExcludeTypes(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls) -> None:
-        cls.dp_original = DeepPot(FROZEN_MODEL)
-        cls.dp_compressed = DeepPot(COMPRESSED_MODEL)
-        cls.coords = np.array(
-            [
-                12.83,
-                2.56,
-                2.18,
-                12.09,
-                2.87,
-                2.74,
-                00.25,
-                3.32,
-                1.68,
-                3.36,
-                3.00,
-                1.81,
-                3.51,
-                2.51,
-                2.60,
-                4.27,
-                3.22,
-                1.56,
-            ]
-        )
-        cls.atype = [0, 1, 1, 0, 1, 1]
-        cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0])
-
-    @classmethod
-    def tearDownClass(cls) -> None:
-        _file_delete(INPUT)
-        _file_delete(FROZEN_MODEL)
-        _file_delete(COMPRESSED_MODEL)
-        _file_delete("out.json")
-        _file_delete("compress.json")
-        _file_delete("checkpoint")
-        _file_delete("model.ckpt.meta")
-        _file_delete("model.ckpt.index")
-        _file_delete("model.ckpt.data-00000-of-00001")
-        _file_delete("model.ckpt-100.meta")
-        _file_delete("model.ckpt-100.index")
-        _file_delete("model.ckpt-100.data-00000-of-00001")
-        _file_delete("model-compression/checkpoint")
-        _file_delete("model-compression/model.ckpt.meta")
-        _file_delete("model-compression/model.ckpt.index")
-        _file_delete("model-compression/model.ckpt.data-00000-of-00001")
-        _file_delete("model-compression")
-        _file_delete("input_v2_compat.json")
-        _file_delete("lcurve.out")
-
-    def test_attrs(self) -> None:
-        self.assertEqual(self.dp_original.get_ntypes(), 2)
-        self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places)
-        self.assertEqual(self.dp_original.get_type_map(), ["O", "H"])
-        self.assertEqual(self.dp_original.get_dim_fparam(), 0)
-        self.assertEqual(self.dp_original.get_dim_aparam(), 0)
-
-        self.assertEqual(self.dp_compressed.get_ntypes(), 2)
-        self.assertAlmostEqual(
-            self.dp_compressed.get_rcut(), 6.0, places=default_places
-        )
-        self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"])
-        self.assertEqual(self.dp_compressed.get_dim_fparam(), 0)
-        self.assertEqual(self.dp_compressed.get_dim_aparam(), 0)
-
-    def test_1frame(self) -> None:
-        ee0, ff0, vv0 = self.dp_original.eval(
-            self.coords, self.box, self.atype, atomic=False
-        )
-        ee1, ff1, vv1 = self.dp_compressed.eval(
-            self.coords, self.box, self.atype, atomic=False
-        )
-        # check shape of the returns
-        nframes = 1
-        natoms = len(self.atype)
-        self.assertEqual(ee0.shape, (nframes, 1))
-        self.assertEqual(ff0.shape, (nframes, natoms, 3))
-        self.assertEqual(vv0.shape, (nframes, 9))
-        self.assertEqual(ee1.shape, (nframes, 1))
-        self.assertEqual(ff1.shape, (nframes, natoms, 3))
-        self.assertEqual(vv1.shape, (nframes, 9))
-        # check values
-        np.testing.assert_almost_equal(ff0, ff1, default_places)
-        np.testing.assert_almost_equal(ee0, ee1, default_places)
-        np.testing.assert_almost_equal(vv0, vv1, default_places)
-
-    def test_1frame_atm(self) -> None:
-        ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(
-            self.coords, self.box, self.atype, atomic=True
-        )
-        ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(
-            self.coords, self.box, self.atype, atomic=True
-        )
-        # check shape of the returns
-        nframes = 1
-        natoms = len(self.atype)
-        self.assertEqual(ee0.shape, (nframes, 1))
-        self.assertEqual(ff0.shape, (nframes, natoms, 3))
-        self.assertEqual(vv0.shape, (nframes, 9))
-        self.assertEqual(ae0.shape, (nframes, natoms, 1))
-        self.assertEqual(av0.shape, (nframes, natoms, 9))
-        self.assertEqual(ee1.shape, (nframes, 1))
-        self.assertEqual(ff1.shape, (nframes, natoms, 3))
-        self.assertEqual(vv1.shape, (nframes, 9))
-        self.assertEqual(ae1.shape, (nframes, natoms, 1))
-        self.assertEqual(av1.shape, (nframes, natoms, 9))
-        # check values
-        np.testing.assert_almost_equal(ff0, ff1, default_places)
-        np.testing.assert_almost_equal(ae0, ae1, default_places)
-        np.testing.assert_almost_equal(av0, av1, default_places)
-        np.testing.assert_almost_equal(ee0, ee1, default_places)
-        np.testing.assert_almost_equal(vv0, vv1, default_places)
-
-    def test_2frame_atm(self) -> None:
-        coords2 = np.concatenate((self.coords, self.coords))
-        box2 = np.concatenate((self.box, self.box))
-        ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(
-            coords2, box2, self.atype, atomic=True
-        )
-        ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(
-            coords2, box2, self.atype, atomic=True
-        )
-        # check shape of the returns
-        nframes = 2
-        natoms = len(self.atype)
-        self.assertEqual(ee0.shape, (nframes, 1))
-        self.assertEqual(ff0.shape, (nframes, natoms, 3))
-        self.assertEqual(vv0.shape, (nframes, 9))
-        self.assertEqual(ae0.shape, (nframes, natoms, 1))
-        self.assertEqual(av0.shape, (nframes, natoms, 9))
-        self.assertEqual(ee1.shape, (nframes, 1))
-        self.assertEqual(ff1.shape, (nframes, natoms, 3))
-        self.assertEqual(vv1.shape, (nframes, 9))
-        self.assertEqual(ae1.shape, (nframes, natoms, 1))
-        self.assertEqual(av1.shape, (nframes, natoms, 9))
-
-        # check values
-        np.testing.assert_almost_equal(ff0, ff1, default_places)
-        np.testing.assert_almost_equal(ae0, ae1, default_places)
-        np.testing.assert_almost_equal(av0, av1, default_places)
-        np.testing.assert_almost_equal(ee0, ee1, default_places)
-        np.testing.assert_almost_equal(vv0, vv1, default_places)
diff --git a/source/tests/tf/test_model_compression_se_t.py b/source/tests/tf/test_model_compression_se_t.py
index 36373d13d5..8b1eb05122 100644
--- a/source/tests/tf/test_model_compression_se_t.py
+++ b/source/tests/tf/test_model_compression_se_t.py
@@ -427,130 +427,6 @@ def test_ase(self) -> None:
         np.testing.assert_almost_equal(ee0, ee1, default_places)
 
 
-class TestDeepPotAPBCExcludeTypes(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls) -> None:
-        cls.dp_original = DeepPot(FROZEN_MODEL)
-        cls.dp_compressed = DeepPot(COMPRESSED_MODEL)
-        cls.coords = np.array(
-            [
-                12.83,
-                2.56,
-                2.18,
-                12.09,
-                2.87,
-                2.74,
-                00.25,
-                3.32,
-                1.68,
-                3.36,
-                3.00,
-                1.81,
-                3.51,
-                2.51,
-                2.60,
-                4.27,
-                3.22,
-                1.56,
-            ]
-        )
-        cls.atype = [0, 1, 1, 0, 1, 1]
-        cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0])
-
-    def test_attrs(self) -> None:
-        self.assertEqual(self.dp_original.get_ntypes(), 2)
-        self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places)
-        self.assertEqual(self.dp_original.get_type_map(), ["O", "H"])
-        self.assertEqual(self.dp_original.get_dim_fparam(), 0)
-        self.assertEqual(self.dp_original.get_dim_aparam(), 0)
-
-        self.assertEqual(self.dp_compressed.get_ntypes(), 2)
-        self.assertAlmostEqual(
-            self.dp_compressed.get_rcut(), 6.0, places=default_places
-        )
-        self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"])
-        self.assertEqual(self.dp_compressed.get_dim_fparam(), 0)
-        self.assertEqual(self.dp_compressed.get_dim_aparam(), 0)
-
-    def test_1frame(self) -> None:
-        ee0, ff0, vv0 = self.dp_original.eval(
-            self.coords, self.box, self.atype, atomic=False
-        )
-        ee1, ff1, vv1 = self.dp_compressed.eval(
-            self.coords, self.box, self.atype, atomic=False
-        )
-        # check shape of the returns
-        nframes = 1
-        natoms = len(self.atype)
-        self.assertEqual(ee0.shape, (nframes, 1))
-        self.assertEqual(ff0.shape, (nframes, natoms, 3))
-        self.assertEqual(vv0.shape, (nframes, 9))
-        self.assertEqual(ee1.shape, (nframes, 1))
-        self.assertEqual(ff1.shape, (nframes, natoms, 3))
-        self.assertEqual(vv1.shape, (nframes, 9))
-        # check values
-        np.testing.assert_almost_equal(ff0, ff1, default_places)
-        np.testing.assert_almost_equal(ee0, ee1, default_places)
-        np.testing.assert_almost_equal(vv0, vv1, default_places)
-
-    def test_1frame_atm(self) -> None:
-        ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(
-            self.coords, self.box, self.atype, atomic=True
-        )
-        ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(
-            self.coords, self.box, self.atype, atomic=True
-        )
-        # check shape of the returns
-        nframes = 1
-        natoms = len(self.atype)
-        self.assertEqual(ee0.shape, (nframes, 1))
-        self.assertEqual(ff0.shape, (nframes, natoms, 3))
-        self.assertEqual(vv0.shape, (nframes, 9))
-        self.assertEqual(ae0.shape, (nframes, natoms, 1))
-        self.assertEqual(av0.shape, (nframes, natoms, 9))
-        self.assertEqual(ee1.shape, (nframes, 1))
-        self.assertEqual(ff1.shape, (nframes, natoms, 3))
-        self.assertEqual(vv1.shape, (nframes, 9))
-        self.assertEqual(ae1.shape, (nframes, natoms, 1))
-        self.assertEqual(av1.shape, (nframes, natoms, 9))
-        # check values
-        np.testing.assert_almost_equal(ff0, ff1, default_places)
-        np.testing.assert_almost_equal(ae0, ae1, default_places)
-        np.testing.assert_almost_equal(av0, av1, default_places)
-        np.testing.assert_almost_equal(ee0, ee1, default_places)
-        np.testing.assert_almost_equal(vv0, vv1, default_places)
-
-    def test_2frame_atm(self) -> None:
-        coords2 = np.concatenate((self.coords, self.coords))
-        box2 = np.concatenate((self.box, self.box))
-        ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(
-            coords2, box2, self.atype, atomic=True
-        )
-        ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(
-            coords2, box2, self.atype, atomic=True
-        )
-        # check shape of the returns
-        nframes = 2
-        natoms = len(self.atype)
-        self.assertEqual(ee0.shape, (nframes, 1))
-        self.assertEqual(ff0.shape, (nframes, natoms, 3))
-        self.assertEqual(vv0.shape, (nframes, 9))
-        self.assertEqual(ae0.shape, (nframes, natoms, 1))
-        self.assertEqual(av0.shape, (nframes, natoms, 9))
-        self.assertEqual(ee1.shape, (nframes, 1))
-        self.assertEqual(ff1.shape, (nframes, natoms, 3))
-        self.assertEqual(vv1.shape, (nframes, 9))
-        self.assertEqual(ae1.shape, (nframes, natoms, 1))
-        self.assertEqual(av1.shape, (nframes, natoms, 9))
-
-        # check values
-        np.testing.assert_almost_equal(ff0, ff1, default_places)
-        np.testing.assert_almost_equal(ae0, ae1, default_places)
-        np.testing.assert_almost_equal(av0, av1, default_places)
-        np.testing.assert_almost_equal(ee0, ee1, default_places)
-        np.testing.assert_almost_equal(vv0, vv1, default_places)
-
-
 class TestDeepPotAPBC2(TestDeepPotAPBC):
     @classmethod
     def setUpClass(cls) -> None:
diff --git a/source/tests/tf/test_multi_device_shape_validation.py b/source/tests/tf/test_multi_device_shape_validation.py
new file mode 100644
index 0000000000..8cf8fa801a
--- /dev/null
+++ b/source/tests/tf/test_multi_device_shape_validation.py
@@ -0,0 +1,391 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Regression tests for flattened TensorFlow custom-op input dimensions."""
+
+import unittest
+
+from deepmd.tf.env import (
+    GLOBAL_TF_FLOAT_PRECISION,
+    op_grads_module,
+    op_module,
+    tf,
+)
+
+
+class TestMultiDeviceShapeValidation(tf.test.TestCase):
+    """Ensure malformed flattened widths fail before native kernel dispatch."""
+
+    def setUp(self) -> None:
+        self.sess = self.cached_session().__enter__()
+        self.nloc = 2
+        self.nnei = 1
+        self.ndescrpt = 4
+        self.natoms = tf.constant([self.nloc, self.nloc, 1], dtype=tf.int32)
+
+    def _floats(self, width: int):
+        """Create one frame of flattened floating-point custom-op input."""
+        return tf.zeros([1, width], dtype=GLOBAL_TF_FLOAT_PRECISION)
+
+    def _nlist(self, width: int):
+        """Create one frame of flattened neighbor indices."""
+        return tf.zeros([1, width], dtype=tf.int32)
+
+    def test_negative_nloc_is_rejected(self) -> None:
+        natoms = tf.constant([-1, 0, 0], dtype=tf.int32)
+
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError,
+            r"number of local atoms should be non-negative",
+        ):
+            self.sess.run(
+                op_module.prod_force_se_a(
+                    self._floats(0),
+                    self._floats(0),
+                    self._nlist(0),
+                    natoms,
+                    n_a_sel=0,
+                    n_r_sel=0,
+                )
+            )
+
+    def test_zero_nloc_rejects_nonempty_flattened_width(self) -> None:
+        natoms = tf.constant([0, 0, 0], dtype=tf.int32)
+
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError,
+            r"net deriv width should be zero when nloc is zero",
+        ):
+            self.sess.run(
+                op_module.prod_force_se_a(
+                    self._floats(1),
+                    self._floats(0),
+                    self._nlist(0),
+                    natoms,
+                    n_a_sel=0,
+                    n_r_sel=0,
+                )
+            )
+
+    def test_zero_nloc_accepts_empty_flattened_widths(self) -> None:
+        natoms = tf.constant([0, 0, 0], dtype=tf.int32)
+
+        result = self.sess.run(
+            op_module.prod_force_se_a(
+                self._floats(0),
+                self._floats(0),
+                self._nlist(0),
+                natoms,
+                n_a_sel=0,
+                n_r_sel=0,
+            )
+        )
+
+        self.assertEqual(result.shape, (1, 0))
+
+    def test_prod_force_rejects_partial_net_deriv_atom(self) -> None:
+        # The old integer division truncated 9 / 2 to four descriptors and
+        # allowed the extra value to survive until raw pointer dispatch.
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError,
+            r"net deriv width 9 should be divisible by nloc 2",
+        ):
+            self.sess.run(
+                op_module.prod_force_se_a(
+                    self._floats(self.nloc * self.ndescrpt + 1),
+                    self._floats(self.nloc * self.ndescrpt * 3),
+                    self._nlist(self.nloc * self.nnei),
+                    self.natoms,
+                    n_a_sel=self.nnei,
+                    n_r_sel=0,
+                )
+            )
+
+    def test_prod_force_rejects_in_deriv_width_mismatch(self) -> None:
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, r"number of descriptors should match"
+        ):
+            self.sess.run(
+                op_module.prod_force_se_a(
+                    self._floats(self.nloc * self.ndescrpt),
+                    self._floats(self.nloc * self.ndescrpt * 3 - 1),
+                    self._nlist(self.nloc * self.nnei),
+                    self.natoms,
+                    n_a_sel=self.nnei,
+                    n_r_sel=0,
+                )
+            )
+
+    def test_prod_force_r_rejects_descriptor_stride_mismatch(self) -> None:
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError,
+            r"descriptor width should equal neighbor width",
+        ):
+            self.sess.run(
+                op_module.prod_force_se_r(
+                    self._floats(self.nloc * (self.nnei + 1)),
+                    self._floats(self.nloc * (self.nnei + 1) * 3),
+                    self._nlist(self.nloc * self.nnei),
+                    self.natoms,
+                )
+            )
+
+    def test_forward_ops_reject_nall_smaller_than_nloc(self) -> None:
+        natoms = tf.constant([self.nloc, self.nloc - 1, 1], dtype=tf.int32)
+        operations = (
+            op_module.prod_force_se_a(
+                self._floats(self.nloc * self.ndescrpt),
+                self._floats(self.nloc * self.ndescrpt * 3),
+                self._nlist(self.nloc * self.nnei),
+                natoms,
+                n_a_sel=self.nnei,
+                n_r_sel=0,
+            ),
+            op_module.prod_force_se_r(
+                self._floats(self.nloc * self.nnei),
+                self._floats(self.nloc * self.nnei * 3),
+                self._nlist(self.nloc * self.nnei),
+                natoms,
+            ),
+            op_module.prod_virial_se_a(
+                self._floats(self.nloc * self.ndescrpt),
+                self._floats(self.nloc * self.ndescrpt * 3),
+                self._floats(self.nloc * self.nnei * 3),
+                self._nlist(self.nloc * self.nnei),
+                natoms,
+                n_a_sel=self.nnei,
+                n_r_sel=0,
+            ),
+            op_module.prod_virial_se_r(
+                self._floats(self.nloc * self.nnei),
+                self._floats(self.nloc * self.nnei * 3),
+                self._floats(self.nloc * self.nnei * 3),
+                self._nlist(self.nloc * self.nnei),
+                natoms,
+            ),
+        )
+        for operation in operations:
+            with (
+                self.subTest(operation=operation),
+                self.assertRaisesRegex(
+                    tf.errors.InvalidArgumentError,
+                    r"number of all atoms should be at least nloc",
+                ),
+            ):
+                self.sess.run(operation)
+
+    def test_prod_virial_rejects_partial_net_deriv_atom(self) -> None:
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError,
+            r"net deriv width 9 should be divisible by nloc 2",
+        ):
+            self.sess.run(
+                op_module.prod_virial_se_a(
+                    self._floats(self.nloc * self.ndescrpt + 1),
+                    self._floats(self.nloc * self.ndescrpt * 3),
+                    self._floats(self.nloc * self.nnei * 3),
+                    self._nlist(self.nloc * self.nnei),
+                    self.natoms,
+                    n_a_sel=self.nnei,
+                    n_r_sel=0,
+                )
+            )
+
+    def test_prod_virial_rejects_descriptor_stride_mismatch(self) -> None:
+        operations = (
+            op_module.prod_virial_se_a(
+                self._floats(self.nloc * self.ndescrpt * 2),
+                self._floats(self.nloc * self.ndescrpt * 2 * 3),
+                self._floats(self.nloc * self.nnei * 3),
+                self._nlist(self.nloc * self.nnei),
+                self.natoms,
+                n_a_sel=self.nnei,
+                n_r_sel=0,
+            ),
+            op_module.prod_virial_se_r(
+                self._floats(self.nloc * (self.nnei + 1)),
+                self._floats(self.nloc * (self.nnei + 1) * 3),
+                self._floats(self.nloc * self.nnei * 3),
+                self._nlist(self.nloc * self.nnei),
+                self.natoms,
+            ),
+        )
+        messages = (
+            r"descriptor width should be four times neighbor width",
+            r"descriptor width should equal neighbor width",
+        )
+        for operation, message in zip(operations, messages, strict=True):
+            with (
+                self.subTest(message=message),
+                self.assertRaisesRegex(tf.errors.InvalidArgumentError, message),
+            ):
+                self.sess.run(operation)
+
+    def test_prod_force_grad_rejects_partial_nlist_atom(self) -> None:
+        # Fixed-width placeholders in the original tests rejected this feed
+        # before the custom op ran, leaving its release-build checks untested.
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError,
+            r"nlist width 3 should be divisible by nloc 2",
+        ):
+            self.sess.run(
+                op_grads_module.prod_force_se_a_grad(
+                    self._floats(self.nloc * 3),
+                    self._floats(self.nloc * self.ndescrpt),
+                    self._floats(self.nloc * self.ndescrpt * 3),
+                    self._nlist(self.nloc * self.nnei + 1),
+                    self.natoms,
+                    n_a_sel=self.nnei,
+                    n_r_sel=0,
+                )
+            )
+
+    def test_prod_force_r_grad_rejects_partial_nlist_atom(self) -> None:
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError,
+            r"nlist width 3 should be divisible by nloc 2",
+        ):
+            self.sess.run(
+                op_grads_module.prod_force_se_r_grad(
+                    self._floats(self.nloc * 3),
+                    self._floats(self.nloc * self.nnei),
+                    self._floats(self.nloc * self.nnei * 3),
+                    self._nlist(self.nloc * self.nnei + 1),
+                    self.natoms,
+                )
+            )
+
+    def test_prod_virial_grad_rejects_descriptor_stride_mismatch(self) -> None:
+        mismatched_ndescrpt = self.ndescrpt * 2
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError,
+            r"descriptor width should be four times neighbor width",
+        ):
+            self.sess.run(
+                op_grads_module.prod_virial_se_a_grad(
+                    self._floats(9),
+                    self._floats(self.nloc * mismatched_ndescrpt),
+                    self._floats(self.nloc * mismatched_ndescrpt * 3),
+                    self._floats(self.nloc * self.nnei * 3),
+                    self._nlist(self.nloc * self.nnei),
+                    self.natoms,
+                    n_a_sel=self.nnei,
+                    n_r_sel=0,
+                )
+            )
+
+    def test_prod_virial_grad_rejects_rij_width_mismatch(self) -> None:
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, r"dim of rij should be  nnei \* 3"
+        ):
+            self.sess.run(
+                op_grads_module.prod_virial_se_a_grad(
+                    self._floats(9),
+                    self._floats(self.nloc * self.ndescrpt),
+                    self._floats(self.nloc * self.ndescrpt * 3),
+                    self._floats(self.nloc * self.nnei * 3 - 1),
+                    self._nlist(self.nloc * self.nnei),
+                    self.natoms,
+                    n_a_sel=self.nnei,
+                    n_r_sel=0,
+                )
+            )
+
+    def test_prod_virial_r_grad_rejects_partial_net_deriv_atom(self) -> None:
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError,
+            r"net deriv width 3 should be divisible by nloc 2",
+        ):
+            self.sess.run(
+                op_grads_module.prod_virial_se_r_grad(
+                    self._floats(9),
+                    self._floats(self.nloc * self.nnei + 1),
+                    self._floats(self.nloc * self.nnei * 3),
+                    self._floats(self.nloc * self.nnei * 3),
+                    self._nlist(self.nloc * self.nnei),
+                    self.natoms,
+                )
+            )
+
+    @unittest.skipUnless(tf.test.is_gpu_available(), "GPU is required")
+    def test_gpu_zero_neighbor_work_returns_zero_outputs(self) -> None:
+        natoms = tf.constant([1, 1, 1], dtype=tf.int32)
+        with tf.device("/GPU:0"):
+            empty_floats = self._floats(0)
+            empty_nlist = self._nlist(0)
+            force_grad = self._floats(3)
+            virial_grad = self._floats(9)
+            force_a = op_module.prod_force_se_a(
+                empty_floats,
+                empty_floats,
+                empty_nlist,
+                natoms,
+                n_a_sel=0,
+                n_r_sel=0,
+            )
+            force_r = op_module.prod_force_se_r(
+                empty_floats, empty_floats, empty_nlist, natoms
+            )
+            force_grad_a = op_grads_module.prod_force_se_a_grad(
+                force_grad,
+                empty_floats,
+                empty_floats,
+                empty_nlist,
+                natoms,
+                n_a_sel=0,
+                n_r_sel=0,
+            )
+            force_grad_r = op_grads_module.prod_force_se_r_grad(
+                force_grad, empty_floats, empty_floats, empty_nlist, natoms
+            )
+            virial_a = op_module.prod_virial_se_a(
+                empty_floats,
+                empty_floats,
+                empty_floats,
+                empty_nlist,
+                natoms,
+                n_a_sel=0,
+                n_r_sel=0,
+            )
+            virial_r = op_module.prod_virial_se_r(
+                empty_floats,
+                empty_floats,
+                empty_floats,
+                empty_nlist,
+                natoms,
+            )
+            virial_grad_a = op_grads_module.prod_virial_se_a_grad(
+                virial_grad,
+                empty_floats,
+                empty_floats,
+                empty_floats,
+                empty_nlist,
+                natoms,
+                n_a_sel=0,
+                n_r_sel=0,
+            )
+            virial_grad_r = op_grads_module.prod_virial_se_r_grad(
+                virial_grad,
+                empty_floats,
+                empty_floats,
+                empty_floats,
+                empty_nlist,
+                natoms,
+            )
+
+        tensors = (
+            force_a,
+            force_r,
+            force_grad_a,
+            force_grad_r,
+            *virial_a,
+            *virial_r,
+            virial_grad_a,
+            virial_grad_r,
+        )
+        for tensor in tensors:
+            self.assertIn("GPU:0", tensor.device)
+            result = self.sess.run(tensor)
+            self.assertAllEqual(result, result * 0)
+
+
+if __name__ == "__main__":
+    tf.test.main()
diff --git a/source/tests/tf/test_tabulate_gpu_size_validation.py b/source/tests/tf/test_tabulate_gpu_size_validation.py
new file mode 100644
index 0000000000..8a856bcced
--- /dev/null
+++ b/source/tests/tf/test_tabulate_gpu_size_validation.py
@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+from deepmd.tf.env import (
+    op_module,
+    tf,
+)
+
+
+@unittest.skipUnless(
+    tf.test.is_gpu_available(), reason="GPU tabulation validation requires a GPU"
+)
+class TestTabulateGpuSizeValidation(unittest.TestCase):
+    """Ensure invalid tabulation widths are rejected before GPU launch."""
+
+    error_message = "last_layer_size must be between 1 and 1024"
+
+    @staticmethod
+    def _build_ops(last_layer_size: int) -> dict[str, object]:
+        """Build all TensorFlow tabulation variants with a common width."""
+        dtype = tf.float64
+        table_width = max(1, 6 * last_layer_size)
+        with tf.device("/CPU:0"):
+            table_info = tf.constant([0.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=dtype)
+
+        with tf.device("/GPU:0"):
+            table = tf.zeros([1, table_width], dtype=dtype)
+            em_x = tf.zeros([1, 1], dtype=dtype)
+
+            em_a = tf.zeros([1, 1, 4], dtype=dtype)
+            descriptor_a = tf.zeros([1, 4, last_layer_size], dtype=dtype)
+            dy_a = tf.zeros_like(descriptor_a)
+            dz_dem_x_a = tf.zeros_like(em_x)
+            dz_dem_a = tf.zeros_like(em_a)
+            two_embed = tf.zeros([1, last_layer_size], dtype=dtype)
+            dz_dtwo = tf.zeros_like(two_embed)
+
+            em_t = tf.zeros([1, 1, 1], dtype=dtype)
+            descriptor_t = tf.zeros([1, last_layer_size], dtype=dtype)
+            dy_t = tf.zeros_like(descriptor_t)
+            dz_dem_x_t = tf.zeros_like(em_x)
+            dz_dem_t = tf.zeros_like(em_t)
+
+            em_r = tf.zeros([1, 1], dtype=dtype)
+            descriptor_r = tf.zeros([1, 1, last_layer_size], dtype=dtype)
+            dy_r = tf.zeros_like(descriptor_r)
+            dz_dem_r = tf.zeros_like(em_r)
+
+            return {
+                "se_a_forward": op_module.tabulate_fusion_se_a(
+                    table,
+                    table_info,
+                    em_x,
+                    em_a,
+                    last_layer_size=last_layer_size,
+                ),
+                "se_a_grad": op_module.tabulate_fusion_se_a_grad(
+                    table, table_info, em_x, em_a, dy_a, descriptor_a
+                ),
+                "se_a_grad_grad": op_module.tabulate_fusion_se_a_grad_grad(
+                    table,
+                    table_info,
+                    em_x,
+                    em_a,
+                    dz_dem_x_a,
+                    dz_dem_a,
+                    descriptor_a,
+                ),
+                "se_atten_forward": op_module.tabulate_fusion_se_atten(
+                    table,
+                    table_info,
+                    em_x,
+                    em_a,
+                    two_embed,
+                    last_layer_size=last_layer_size,
+                ),
+                "se_atten_grad": op_module.tabulate_fusion_se_atten_grad(
+                    table,
+                    table_info,
+                    em_x,
+                    em_a,
+                    two_embed,
+                    dy_a,
+                    descriptor_a,
+                ),
+                "se_atten_grad_grad": (
+                    op_module.tabulate_fusion_se_atten_grad_grad(
+                        table,
+                        table_info,
+                        em_x,
+                        em_a,
+                        two_embed,
+                        dz_dem_x_a,
+                        dz_dem_a,
+                        dz_dtwo,
+                        descriptor_a,
+                    )
+                ),
+                "se_t_forward": op_module.tabulate_fusion_se_t(
+                    table,
+                    table_info,
+                    em_x,
+                    em_t,
+                    last_layer_size=last_layer_size,
+                ),
+                "se_t_grad": op_module.tabulate_fusion_se_t_grad(
+                    table, table_info, em_x, em_t, dy_t, descriptor_t
+                ),
+                "se_t_grad_grad": op_module.tabulate_fusion_se_t_grad_grad(
+                    table,
+                    table_info,
+                    em_x,
+                    em_t,
+                    dz_dem_x_t,
+                    dz_dem_t,
+                    descriptor_t,
+                ),
+                "se_r_forward": op_module.tabulate_fusion_se_r(
+                    table,
+                    table_info,
+                    em_r,
+                    last_layer_size=last_layer_size,
+                ),
+                "se_r_grad": op_module.tabulate_fusion_se_r_grad(
+                    table, table_info, em_r, dy_r, descriptor_r
+                ),
+                "se_r_grad_grad": op_module.tabulate_fusion_se_r_grad_grad(
+                    table, table_info, em_r, dz_dem_r, descriptor_r
+                ),
+            }
+
+    def _assert_invalid(self, last_layer_size: int, names: tuple[str, ...]) -> None:
+        graph = tf.Graph()
+        with graph.as_default():
+            ops = self._build_ops(last_layer_size)
+
+        config = tf.ConfigProto(allow_soft_placement=False)
+        # Constant folding can execute an all-constant custom op on the CPU
+        # before placement, bypassing the GPU-only launch validation.
+        config.graph_options.optimizer_options.opt_level = -1
+        config.graph_options.rewrite_options.disable_meta_optimizer = True
+        config.gpu_options.allow_growth = True
+        with tf.Session(graph=graph, config=config) as sess:
+            for name in names:
+                with self.subTest(op=name, last_layer_size=last_layer_size):
+                    with self.assertRaisesRegex(
+                        tf.errors.InvalidArgumentError, self.error_message
+                    ):
+                        sess.run(ops[name])
+
+    def test_rejects_oversized_width_for_all_gpu_paths(self) -> None:
+        """Cover every forward, first-gradient, and grad-grad GPU wrapper."""
+        names = tuple(
+            f"{descriptor}_{stage}"
+            for descriptor in ("se_a", "se_atten", "se_t", "se_r")
+            for stage in ("forward", "grad", "grad_grad")
+        )
+        self._assert_invalid(1025, names)
+
+    def test_rejects_zero_width(self) -> None:
+        """Cover attribute-derived and descriptor-derived zero widths."""
+        self._assert_invalid(0, ("se_a_forward", "se_t_grad"))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/tf/test_train.py b/source/tests/tf/test_train.py
index 4a0fb55494..f54b66aaad 100644
--- a/source/tests/tf/test_train.py
+++ b/source/tests/tf/test_train.py
@@ -158,34 +158,6 @@ def test_update_sel_atten_int(self, sel_mock, get_data_mock) -> None:
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    @patch("deepmd.tf.entrypoints.train.get_data")
-    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat")
-    def test_update_sel_atten_list(self, sel_mock, get_data_mock) -> None:
-        sel_mock.return_value = self.mock_min_nbor_dist, [25]
-        get_data_mock.return_value = None
-        jdata = {
-            "model": {
-                "descriptor": {
-                    "type": "se_atten",
-                    "sel": 30,
-                    "rcut": 6,
-                }
-            },
-            "training": {"training_data": {}},
-        }
-        expected_out = {
-            "model": {
-                "descriptor": {
-                    "type": "se_atten",
-                    "sel": 30,
-                    "rcut": 6,
-                }
-            },
-            "training": {"training_data": {}},
-        }
-        jdata = update_sel(jdata)
-        self.assertEqual(jdata, expected_out)
-
     @patch("deepmd.tf.entrypoints.train.get_data")
     def test_skip_loc_frame(self, get_data_mock) -> None:
         get_data_mock.return_value = None
diff --git a/source/tests/universal/common/cases/atomic_model/utils.py b/source/tests/universal/common/cases/atomic_model/utils.py
index fbf6600320..d321b45e96 100644
--- a/source/tests/universal/common/cases/atomic_model/utils.py
+++ b/source/tests/universal/common/cases/atomic_model/utils.py
@@ -49,50 +49,17 @@ class AtomicModelTestCase:
     epsilon_dict: dict[str, float | None]
     """Dictionary of epsilons in each test."""
 
-    def test_get_type_map(self) -> None:
-        """Test get_type_map."""
+    def test_metadata(self) -> None:
+        """Model metadata survives serialization and supported JIT conversion."""
         for module in self.modules_to_test:
             self.assertEqual(module.get_type_map(), self.expected_type_map)
-
-    def test_get_rcut(self) -> None:
-        """Test get_rcut."""
-        for module in self.modules_to_test:
             self.assertAlmostEqual(module.get_rcut(), self.expected_rcut)
-
-    def test_get_dim_fparam(self) -> None:
-        """Test get_dim_fparam."""
-        for module in self.modules_to_test:
             self.assertEqual(module.get_dim_fparam(), self.expected_dim_fparam)
-
-    def test_get_dim_aparam(self) -> None:
-        """Test get_dim_aparam."""
-        for module in self.modules_to_test:
             self.assertEqual(module.get_dim_aparam(), self.expected_dim_aparam)
-
-    def test_get_sel_type(self) -> None:
-        """Test get_sel_type."""
-        for module in self.modules_to_test:
             self.assertEqual(module.get_sel_type(), self.expected_sel_type)
-
-    def test_is_aparam_nall(self) -> None:
-        """Test is_aparam_nall."""
-        for module in self.modules_to_test:
             self.assertEqual(module.is_aparam_nall(), self.expected_aparam_nall)
-
-    def test_get_nnei(self) -> None:
-        """Test get_nnei."""
-        expected_nnei = sum(self.expected_sel)
-        for module in self.modules_to_test:
-            self.assertEqual(module.get_nnei(), expected_nnei)
-
-    def test_get_ntypes(self) -> None:
-        """Test get_ntypes."""
-        for module in self.modules_to_test:
+            self.assertEqual(module.get_nnei(), sum(self.expected_sel))
             self.assertEqual(module.get_ntypes(), len(self.expected_type_map))
-
-    def test_has_message_passing(self) -> None:
-        """Test has_message_passing."""
-        for module in self.modules_to_test:
             self.assertEqual(
                 module.has_message_passing(), self.expected_has_message_passing
             )
diff --git a/source/tests/universal/common/cases/cases.py b/source/tests/universal/common/cases/cases.py
index 421cbf4556..b8068d1163 100644
--- a/source/tests/universal/common/cases/cases.py
+++ b/source/tests/universal/common/cases/cases.py
@@ -2,6 +2,42 @@
 import numpy as np
 
 
+def type_map_cases() -> list[tuple[tuple[str, ...], tuple[str, ...], bool]]:
+    """Cover one-hot residual widths and electronic-configuration map changes.
+
+    One-hot inputs cross the embedding widths at 8 and 16 types. Electronic
+    configurations have a fixed input width, so permutation, shrink and growth
+    cover their distinct map operations without repeating the width matrix.
+    """
+    elements = (
+        "H",
+        "He",
+        "Li",
+        "Be",
+        "B",
+        "C",
+        "N",
+        "O",
+        "F",
+        "Ne",
+        "Na",
+        "Mg",
+        "Al",
+        "Si",
+        "P",
+        "S",
+        "Cl",
+        "Ar",
+    )
+    small = ("H", "O")
+    maps = (elements, elements[:16], elements[:8], small)
+    return [(old, new, False) for old in maps for new in maps] + [
+        (elements, elements, True),
+        (elements, small, True),
+        (small, elements, True),
+    ]
+
+
 # originally copied from source/tests/pt/model/test_env_mat.py
 class TestCaseSingleFrameWithNlist:
     def setUp(self) -> None:
diff --git a/source/tests/universal/common/cases/descriptor/utils.py b/source/tests/universal/common/cases/descriptor/utils.py
index ca5a750e1c..21da23d11b 100644
--- a/source/tests/universal/common/cases/descriptor/utils.py
+++ b/source/tests/universal/common/cases/descriptor/utils.py
@@ -18,6 +18,7 @@
 )
 from ..cases import (
     TestCaseSingleFrameWithNlist,
+    type_map_cases,
 )
 
 
@@ -144,50 +145,13 @@ def test_change_type_map(self) -> None:
         mapping_device = self.mapping
         # type_map for data and exclude_types
         original_type_map = ["O", "H"]
-        full_type_map_test = [
-            "H",
-            "He",
-            "Li",
-            "Be",
-            "B",
-            "C",
-            "N",
-            "O",
-            "F",
-            "Ne",
-            "Na",
-            "Mg",
-            "Al",
-            "Si",
-            "P",
-            "S",
-            "Cl",
-            "Ar",
-        ]  # 18 elements
         rng = np.random.default_rng(GLOBAL_SEED)
-        for old_tm, new_tm, em, econf in itertools.product(
-            [
-                full_type_map_test[:],  # 18 elements
-                full_type_map_test[
-                    :16
-                ],  # 16 elements, double of tebd default first dim
-                full_type_map_test[:8],  # 8 elements, tebd default first dim
-                ["H", "O"],  # slimmed types
-            ],  # old_type_map
-            [
-                full_type_map_test[:],  # 18 elements
-                full_type_map_test[
-                    :16
-                ],  # 16 elements, double of tebd default first dim
-                full_type_map_test[:8],  # 8 elements, tebd default first dim
-                ["H", "O"],  # slimmed types
-            ],  # new_type_map
-            [[], [[0, 1]], [[1, 1]]],  # exclude_types for original_type_map
-            [False, True],  # use_econf_tebd
+        for (old_tm, new_tm, econf), em in itertools.product(
+            type_map_cases(), [[], [[0, 1]], [[1, 1]]]
         ):
             # use shuffled type_map
-            rng.shuffle(old_tm)
-            rng.shuffle(new_tm)
+            old_tm = rng.permutation(old_tm).tolist()
+            new_tm = rng.permutation(new_tm).tolist()
             old_tm_index = np.array(
                 [old_tm.index(i) for i in original_type_map], dtype=np.int32
             )
diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py
index c51f9aefd1..127df75ecc 100644
--- a/source/tests/universal/common/cases/model/utils.py
+++ b/source/tests/universal/common/cases/model/utils.py
@@ -62,85 +62,32 @@ class ModelTestCase:
     epsilon_dict: dict[str, float | None]
     """Dictionary of epsilons in each test."""
 
-    def test_get_type_map(self) -> None:
-        """Test get_type_map."""
+    def test_metadata(self) -> None:
+        """Model metadata survives serialization and supported JIT conversion."""
         for module in self.modules_to_test:
             self.assertEqual(module.get_type_map(), self.expected_type_map)
-
-    def test_get_rcut(self) -> None:
-        """Test get_rcut."""
-        for module in self.modules_to_test:
             self.assertAlmostEqual(module.get_rcut(), self.expected_rcut)
-
-    def test_get_dim_fparam(self) -> None:
-        """Test get_dim_fparam."""
-        for module in self.modules_to_test:
             self.assertEqual(module.get_dim_fparam(), self.expected_dim_fparam)
-
-    def test_get_dim_aparam(self) -> None:
-        """Test get_dim_aparam."""
-        for module in self.modules_to_test:
             self.assertEqual(module.get_dim_aparam(), self.expected_dim_aparam)
-
-    def test_get_sel_type(self) -> None:
-        """Test get_sel_type."""
-        for module in self.modules_to_test:
             self.assertEqual(module.get_sel_type(), self.expected_sel_type)
-
-    def test_is_aparam_nall(self) -> None:
-        """Test is_aparam_nall."""
-        for module in self.modules_to_test:
             self.assertEqual(module.is_aparam_nall(), self.expected_aparam_nall)
-
-    def test_model_output_type(self) -> None:
-        """Test model_output_type."""
-        for module in self.modules_to_test:
             self.assertEqual(
                 module.model_output_type(), self.expected_model_output_type
             )
-
-    def test_get_nnei(self) -> None:
-        """Test get_nnei."""
-        expected_nnei = sum(self.expected_sel)
-        for module in self.modules_to_test:
-            self.assertEqual(module.get_nnei(), expected_nnei)
-
-    def test_get_ntypes(self) -> None:
-        """Test get_ntypes."""
-        for module in self.modules_to_test:
+            self.assertEqual(module.get_nnei(), sum(self.expected_sel))
             self.assertEqual(module.get_ntypes(), len(self.expected_type_map))
-
-    def test_has_message_passing(self) -> None:
-        """Test has_message_passing."""
-        for module in self.modules_to_test:
             self.assertEqual(
                 module.has_message_passing(), self.expected_has_message_passing
             )
 
-    def test_has_spin(self) -> None:
-        """Test has_spin.
+    def test_capabilities(self) -> None:
+        """Check Python capabilities on the unfrozen model.
 
-        Declared on the base model with a concrete default of ``False``
-        (direct method call, never a ``getattr`` probe); spin model wrappers
-        override it to ``True`` -- which is exactly the suite's ``test_spin``
-        flag. Checked on the raw module only (not the jit-scripted copies in
-        ``modules_to_test``): the capability contract is a Python-API
-        contract, and pt's TorchScript models do not ``@torch.jit.export``
-        it.
+        These methods are not all exported by TorchScript.
         """
         expected = getattr(self, "test_spin", False)
         self.assertEqual(self.module.has_spin(), expected)
 
-    def test_chg_spin_capability_contract(self) -> None:
-        """chg-spin queries are declared on the base model with concrete
-        defaults (False/0/None) -- direct calls, never ``hasattr`` probes.
-
-        ``has_default_chg_spin`` was merged into ``get_default_chg_spin``
-        (predicate: ``get_default_chg_spin() is not None``). Its absence on
-        the dpmodel side is pinned in
-        source/tests/universal/dpmodel/model/test_model.py -- pt is frozen
-        and still declares the (now-redundant) method.
-        """
         assert isinstance(self.module.has_chg_spin_ebd(), bool)
         assert isinstance(self.module.get_dim_chg_spin(), int)
         dcs = self.module.get_default_chg_spin()
@@ -148,12 +95,6 @@ def test_chg_spin_capability_contract(self) -> None:
         # torch.Tensor -- pin the shape contract, not the container type.
         assert dcs is None or len(dcs) == self.module.get_dim_chg_spin()
 
-    def test_property_capability_contract(self) -> None:
-        """Property queries are declared on the base model with concrete
-        defaults: ``get_var_name`` returns ``None`` for non-property models
-        (the support predicate), ``get_intensive`` defaults ``False``,
-        ``get_task_dim`` raises for non-property models.
-        """
         vn = self.module.get_var_name()
         assert vn is None or isinstance(vn, str)
         assert isinstance(self.module.get_intensive(), bool)
diff --git a/source/tests/universal/common/cases/utils/utils.py b/source/tests/universal/common/cases/utils/utils.py
index f9fc697cf2..a6fc110a32 100644
--- a/source/tests/universal/common/cases/utils/utils.py
+++ b/source/tests/universal/common/cases/utils/utils.py
@@ -11,6 +11,7 @@
 )
 from ..cases import (
     TestCaseSingleFrameWithNlist,
+    type_map_cases,
 )
 
 
@@ -31,47 +32,9 @@ def test_change_type_map(self) -> None:
         atype_ext_device = self.atype_ext
         # type_map for data and exclude_types
         original_type_map = ["O", "H"]
-        full_type_map_test = [
-            "H",
-            "He",
-            "Li",
-            "Be",
-            "B",
-            "C",
-            "N",
-            "O",
-            "F",
-            "Ne",
-            "Na",
-            "Mg",
-            "Al",
-            "Si",
-            "P",
-            "S",
-            "Cl",
-            "Ar",
-        ]  # 18 elements
         rng = np.random.default_rng(GLOBAL_SEED)
-        for old_tm, new_tm, neuron, act, econf in itertools.product(
-            [
-                full_type_map_test[:],  # 18 elements
-                full_type_map_test[
-                    :16
-                ],  # 16 elements, double of tebd default first dim
-                full_type_map_test[:8],  # 8 elements, tebd default first dim
-                ["H", "O"],  # slimmed types
-            ],  # old_type_map
-            [
-                full_type_map_test[:],  # 18 elements
-                full_type_map_test[
-                    :16
-                ],  # 16 elements, double of tebd default first dim
-                full_type_map_test[:8],  # 8 elements, tebd default first dim
-                ["H", "O"],  # slimmed types
-            ],  # new_type_map
-            [[8], [8, 16, 32]],  # neuron
-            ["Linear", "tanh"],  # activation_function
-            [False, True],  # use_econf_tebd
+        for (old_tm, new_tm, econf), neuron, act in itertools.product(
+            type_map_cases(), [[8], [8, 16, 32]], ["Linear", "tanh"]
         ):
             do_resnet = neuron[0] in [
                 len(old_tm),
@@ -83,8 +46,8 @@ def test_change_type_map(self) -> None:
                 # `activation_function` must be "Linear" when performing type changing on resnet structure
                 continue
             # use shuffled type_map
-            rng.shuffle(old_tm)
-            rng.shuffle(new_tm)
+            old_tm = rng.permutation(old_tm).tolist()
+            new_tm = rng.permutation(new_tm).tolist()
             old_tm_index = np.array(
                 [old_tm.index(i) for i in original_type_map], dtype=np.int32
             )
diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py
index 85faf46ed9..c6cd002fcd 100644
--- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py
+++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py
@@ -1,8 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import unittest
-from collections import (
-    OrderedDict,
-)
 from typing import (
     Any,
 )
@@ -37,7 +34,6 @@
 )
 
 from ....consistent.common import (
-    parameterize_func,
     parameterized,
 )
 from ....seed import (
@@ -86,22 +82,6 @@ def DescriptorParamSeA(
     return input_dict
 
 
-DescriptorParamSeAList = parameterize_func(
-    DescriptorParamSeA,
-    OrderedDict(
-        {
-            "resnet_dt": (False, True),
-            "type_one_side": (True, False),
-            "exclude_types": ([], [[0, 1]]),
-            "env_protection": (0.0, 1e-8, 1e-2),
-            "precision": ("float64",),
-        }
-    ),
-)
-# to get name for the default function
-DescriptorParamSeA = DescriptorParamSeAList[0]
-
-
 def DescriptorParamSeR(
     ntypes,
     rcut,
@@ -130,22 +110,6 @@ def DescriptorParamSeR(
     return input_dict
 
 
-DescriptorParamSeRList = parameterize_func(
-    DescriptorParamSeR,
-    OrderedDict(
-        {
-            "resnet_dt": (False, True),
-            "type_one_side": (True,),  # type_one_side == False not implemented
-            "exclude_types": ([], [[0, 1]]),
-            "env_protection": (0.0, 1e-8),
-            "precision": ("float64",),
-        }
-    ),
-)
-# to get name for the default function
-DescriptorParamSeR = DescriptorParamSeRList[0]
-
-
 def DescriptorParamSeT(
     ntypes,
     rcut,
@@ -172,21 +136,6 @@ def DescriptorParamSeT(
     return input_dict
 
 
-DescriptorParamSeTList = parameterize_func(
-    DescriptorParamSeT,
-    OrderedDict(
-        {
-            "resnet_dt": (False, True),
-            "exclude_types": ([], [[0, 1]]),
-            "env_protection": (0.0, 1e-8),
-            "precision": ("float64",),
-        }
-    ),
-)
-# to get name for the default function
-DescriptorParamSeT = DescriptorParamSeTList[0]
-
-
 def DescriptorParamSeTTebd(
     ntypes,
     rcut,
@@ -199,7 +148,7 @@ def DescriptorParamSeTTebd(
     tebd_input_mode="concat",
     concat_output_tebd=True,
     resnet_dt=True,
-    set_davg_zero=True,
+    set_davg_zero=False,
     smooth=True,
     use_econf_tebd=False,
     use_tebd_bias=False,
@@ -227,28 +176,6 @@ def DescriptorParamSeTTebd(
     return input_dict
 
 
-DescriptorParamSeTTebdList = parameterize_func(
-    DescriptorParamSeTTebd,
-    OrderedDict(
-        {
-            "tebd_dim": (4,),
-            "tebd_input_mode": ("concat", "strip"),
-            "resnet_dt": (True,),
-            "exclude_types": ([], [[0, 1]]),
-            "env_protection": (0.0,),
-            "set_davg_zero": (False,),
-            "smooth": (True, False),
-            "concat_output_tebd": (True,),
-            "use_econf_tebd": (False, True),
-            "use_tebd_bias": (False,),
-            "precision": ("float64",),
-        }
-    ),
-)
-# to get name for the default function
-DescriptorParamSeTTebd = DescriptorParamSeTTebdList[0]
-
-
 def DescriptorParamDPA1(
     ntypes,
     rcut,
@@ -260,7 +187,7 @@ def DescriptorParamDPA1(
     tebd_dim=4,
     tebd_input_mode="concat",
     attn=20,
-    attn_layer=2,
+    attn_layer=0,
     attn_dotr=True,
     scaling_factor=1.0,
     normalize=True,
@@ -268,8 +195,8 @@ def DescriptorParamDPA1(
     ln_eps=1e-5,
     concat_output_tebd=True,
     resnet_dt=True,
-    type_one_side=True,
-    set_davg_zero=True,
+    type_one_side=False,
+    set_davg_zero=False,
     smooth_type_embedding=True,
     use_econf_tebd=False,
     use_tebd_bias=False,
@@ -306,36 +233,6 @@ def DescriptorParamDPA1(
     return input_dict
 
 
-DescriptorParamDPA1List = parameterize_func(
-    DescriptorParamDPA1,
-    OrderedDict(
-        {
-            "tebd_dim": (4,),
-            "tebd_input_mode": ("concat", "strip"),
-            "resnet_dt": (True,),
-            "type_one_side": (False,),
-            "attn": (20,),
-            "attn_layer": (0, 2),
-            "attn_dotr": (True,),
-            "exclude_types": ([], [[0, 1]]),
-            "env_protection": (0.0,),
-            "set_davg_zero": (False,),
-            "scaling_factor": (1.0,),
-            "normalize": (True,),
-            "temperature": (None, 1.0),
-            "ln_eps": (1e-5,),
-            "smooth_type_embedding": (True, False),
-            "concat_output_tebd": (True,),
-            "use_econf_tebd": (False, True),
-            "use_tebd_bias": (False,),
-            "precision": ("float64",),
-        }
-    ),
-)
-# to get name for the default function
-DescriptorParamDPA1 = DescriptorParamDPA1List[0]
-
-
 def DescriptorParamDPA2(
     ntypes,
     rcut,
@@ -343,9 +240,9 @@ def DescriptorParamDPA2(
     sel,
     type_map,
     repinit_tebd_input_mode="concat",
-    repinit_set_davg_zero=False,
+    repinit_set_davg_zero=True,
     repinit_type_one_side=False,
-    repinit_use_three_body=False,
+    repinit_use_three_body=True,
     repformer_direct_dist=False,
     repformer_update_g1_has_conv=True,
     repformer_update_g1_has_drrd=True,
@@ -357,12 +254,12 @@ def DescriptorParamDPA2(
     repformer_attn2_has_gate=True,
     repformer_update_style="res_avg",
     repformer_update_residual_init="norm",
-    repformer_set_davg_zero=False,
+    repformer_set_davg_zero=True,
     repformer_trainable_ln=True,
     repformer_ln_eps=1e-5,
-    repformer_use_sqrt_nnei=False,
-    repformer_g1_out_conv=False,
-    repformer_g1_out_mlp=False,
+    repformer_use_sqrt_nnei=True,
+    repformer_g1_out_conv=True,
+    repformer_g1_out_mlp=True,
     smooth=True,
     add_tebd_to_repinit_out=True,
     use_econf_tebd=False,
@@ -443,44 +340,6 @@ def DescriptorParamDPA2(
     return input_dict
 
 
-DescriptorParamDPA2List = parameterize_func(
-    DescriptorParamDPA2,
-    OrderedDict(
-        {
-            "repinit_tebd_input_mode": ("concat", "strip"),
-            "repinit_set_davg_zero": (True,),
-            "repinit_type_one_side": (False,),
-            "repinit_use_three_body": (True, False),
-            "repformer_direct_dist": (False,),
-            "repformer_update_g1_has_conv": (True,),
-            "repformer_update_g1_has_drrd": (True,),
-            "repformer_update_g1_has_grrg": (True,),
-            "repformer_update_g1_has_attn": (True,),
-            "repformer_update_g2_has_g1g1": (True,),
-            "repformer_update_g2_has_attn": (True,),
-            "repformer_update_h2": (False,),
-            "repformer_attn2_has_gate": (True,),
-            "repformer_update_style": ("res_avg", "res_residual"),
-            "repformer_update_residual_init": ("norm",),
-            "repformer_set_davg_zero": (True,),
-            "repformer_trainable_ln": (True,),
-            "repformer_ln_eps": (1e-5,),
-            "repformer_use_sqrt_nnei": (True,),
-            "repformer_g1_out_conv": (True,),
-            "repformer_g1_out_mlp": (True,),
-            "smooth": (True, False),
-            "exclude_types": ([], [[0, 1]]),
-            "precision": ("float64",),
-            "add_tebd_to_repinit_out": (True, False),
-            "use_econf_tebd": (False,),
-            "use_tebd_bias": (False,),
-        }
-    ),
-)
-# to get name for the default function
-DescriptorParamDPA2 = DescriptorParamDPA2List[0]
-
-
 def DescriptorParamDPA3(
     ntypes,
     rcut,
@@ -494,15 +353,15 @@ def DescriptorParamDPA3(
     update_residual_init="const",
     update_angle=True,
     n_multi_edge_message=1,
-    a_compress_rate=0,
-    a_compress_e_rate=1,
-    a_compress_use_split=False,
+    a_compress_rate=1,
+    a_compress_e_rate=2,
+    a_compress_use_split=True,
     optim_update=True,
-    smooth_edge_update=False,
-    edge_init_use_dist=False,
-    use_exp_switch=False,
+    smooth_edge_update=True,
+    edge_init_use_dist=True,
+    use_exp_switch=True,
     fix_stat_std=0.3,
-    use_dynamic_sel=False,
+    use_dynamic_sel=True,
     precision="float64",
     use_loc_mapping=True,
     add_chg_spin_ebd=False,
@@ -562,35 +421,6 @@ def DescriptorParamDPA3(
     return input_dict
 
 
-DescriptorParamDPA3List = parameterize_func(
-    DescriptorParamDPA3,
-    OrderedDict(
-        {
-            "update_residual_init": ("const",),
-            "exclude_types": ([], [[0, 1]]),
-            "update_angle": (True,),
-            "a_compress_rate": (1,),
-            "a_compress_e_rate": (2,),
-            "a_compress_use_split": (True,),
-            "optim_update": (True, False),
-            "smooth_edge_update": (True,),
-            "edge_init_use_dist": (True, False),
-            "use_exp_switch": (True, False),
-            "fix_stat_std": (0.3,),
-            "n_multi_edge_message": (1,),
-            "use_dynamic_sel": (True, False),
-            "env_protection": (0.0, 1e-8),
-            "precision": ("float64",),
-            "use_loc_mapping": (True, False),
-        }
-    ),
-)
-
-
-# to get name for the default function
-DescriptorParamDPA3 = DescriptorParamDPA3List[0]
-
-
 def _descriptor_param_variant(param_func, name: str, **fixed_kwargs: Any):
     def wrapper(*args, **kwargs):
         return param_func(*args, **{**kwargs, **fixed_kwargs})
@@ -600,9 +430,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-# Curated descriptor variants for model/atomic-model integration tests. The
-# descriptor tests above still cover the full parameter matrices; model-level
-# tests only need representative toggles to avoid a Cartesian-product blow-up.
+# Descriptor variants exercise independent options in model integration tests.
 DescriptorParamSeAEnergyModelList = (
     # Baseline coverage.
     DescriptorParamSeA,