fix(pt_expt): reuse the stored min_nbor_dist and batch the neighbor statistics

deepmd/main.py

@@ -668,6 +668,15 @@ def main_parser() -> argparse.ArgumentParser:
         default=None,
         help="The training script of the input frozen model",
     )
+    parser_compress.add_argument(
+        "--recompute-min-nbor-dist",
+        action="store_true",
+        help="(Supported backend: PyTorch Exportable) Ignore the minimal neighbor "
+        "distance saved in the model and recompute it from the training data. "
+        "Requires -t,--training-script. Only applies to models whose compression "
+        "consumes the statistic; it is ignored for a model that tabulates over an "
+        "analytically bounded domain",
+    )
     parser_compress.add_argument(
         "--head",
         "--model-branch",

deepmd/pt_expt/entrypoints/compress.py

@@ -2,6 +2,9 @@
 """Compress a pt_expt model (.pte) by tabulating embedding nets."""
 
 import logging
+from typing import (
+    Any,
+)
 
 from deepmd.pt_expt.utils.serialization import (
     deserialize_to_file,
@@ -11,13 +14,75 @@
 log = logging.getLogger(__name__)
 
 
+def _read_saved_min_nbor_dist(model: Any, model_dict: dict) -> tuple[float | None, str]:
+    """Read the stored minimal neighbor distance and where it was read from.
+
+    ``@variables`` is the cross-backend location of this value: it is written
+    by :mod:`deepmd.pt.utils.serialization` and read back by the PyTorch and
+    Paddle backends, so a ``.pt2`` produced by ``dp convert-backend`` carries
+    the value there rather than inside the serialized model dict.
+
+    Returns
+    -------
+    float or None
+        The stored minimal neighbor distance, None if the model has none.
+    str
+        Human-readable description of where the value was read from.
+    """
+    min_nbor_dist = model.get_min_nbor_dist()
+    if min_nbor_dist is not None:
+        return float(min_nbor_dist), "the model"
+    min_nbor_dist = model_dict.get("min_nbor_dist")
+    if min_nbor_dist is not None:
+        return float(min_nbor_dist), "the model file"
+    min_nbor_dist = (model_dict.get("@variables") or {}).get("min_nbor_dist")
+    if min_nbor_dist is not None:
+        return float(min_nbor_dist), "the model file (@variables)"
+    return None, ""
+
+
+def _measure_min_nbor_dist(training_script: str) -> float:
+    """Measure the minimal neighbor distance from the training data.
+
+    This is a dense all-pairs pass over every frame, so callers run it only
+    when the tabulation actually consumes the statistic.
+
+    Returns
+    -------
+    float
+        The shortest neighbor distance in the training data.
+    """
+    from deepmd.common import (
+        j_loader,
+    )
+    from deepmd.pt_expt.utils.update_sel import (
+        UpdateSel,
+    )
+    from deepmd.utils.compat import (
+        update_deepmd_input,
+    )
+    from deepmd.utils.data_system import (
+        get_data,
+    )
+
+    jdata = update_deepmd_input(j_loader(training_script))
+    train_data = get_data(
+        jdata["training"]["training_data"],
+        0,
+        jdata["model"].get("type_map", None),
+        None,
+    )
+    return UpdateSel().get_min_nbor_dist(train_data)
+
+
 def enable_compression(
     input_file: str,
     output: str,
     stride: float = 0.01,
     extrapolate: int = 5,
     check_frequency: int = -1,
     training_script: str | None = None,
+    recompute_min_nbor_dist: bool = False,
 ) -> None:
     """Compress a .pte model by tabulating embedding nets.
 
@@ -36,6 +101,12 @@ def enable_compression(
     training_script : str or None
         Path to training script, used to compute min_nbor_dist if not
         stored in the model.
+    recompute_min_nbor_dist : bool
+        Ignore the min_nbor_dist stored in the model and recompute it from
+        the training data. Requires training_script. Only applies to models
+        whose compression consumes the statistic; it is ignored, with a
+        warning, for a model that tabulates over an analytically bounded
+        domain.
     """
     from deepmd.pt_expt.model.model import (
         BaseModel,
@@ -45,55 +116,59 @@ def enable_compression(
     model_dict = serialize_from_file(input_file)
     model = BaseModel.deserialize(model_dict["model"])
 
-    # 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")
-    if min_nbor_dist is None and not model.compression_needs_min_nbor_dist():
-        log.info(
-            "The model tabulates over an analytically bounded domain; "
-            "skipping the neighbor statistics."
-        )
-        min_nbor_dist = 0.0
-    if min_nbor_dist is None:
-        log.info(
-            "Minimal neighbor distance is not saved in the model, "
-            "compute it from the training data."
-        )
-        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)."
+    # 2. Resolve min_nbor_dist, branching on whether compression consumes it.
+    #    Measuring it is a dense all-pairs pass over the training data, so a
+    #    model whose tabulation domain is fixed analytically never scans.
+    saved_min_nbor_dist, source = _read_saved_min_nbor_dist(model, model_dict)
+    if model.compression_needs_min_nbor_dist():
+        if recompute_min_nbor_dist:
+            log.info(
+                "Recompute the minimal neighbor distance from the training data, "
+                "ignoring the one saved in the model."
             )
-        from deepmd.common import (
-            j_loader,
-        )
-        from deepmd.pt_expt.utils.update_sel import (
-            UpdateSel,
-        )
-        from deepmd.utils.compat import (
-            update_deepmd_input,
-        )
-        from deepmd.utils.data_system import (
-            get_data,
-        )
-
-        jdata = j_loader(training_script)
-        jdata = update_deepmd_input(jdata)
-        type_map = jdata["model"].get("type_map", None)
-        train_data = get_data(
-            jdata["training"]["training_data"],
-            0,
-            type_map,
-            None,
-        )
-        update_sel = UpdateSel()
-        min_nbor_dist = update_sel.get_min_nbor_dist(train_data)
+            if training_script is None:
+                raise ValueError(
+                    "Recomputing the minimal neighbor distance requires the "
+                    "training script and data (via -t,--training-script)."
+                )
+            min_nbor_dist = _measure_min_nbor_dist(training_script)
+        elif saved_min_nbor_dist is not None:
+            min_nbor_dist = saved_min_nbor_dist
+            log.info(f"Minimal neighbor distance read from {source}: {min_nbor_dist:f}")
+        else:
+            log.info(
+                "Minimal neighbor distance is not saved in the model, "
+                "compute it from the training data."
+            )
+            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)."
+                )
+            min_nbor_dist = _measure_min_nbor_dist(training_script)
+    else:
+        if recompute_min_nbor_dist:
+            log.warning(
+                "--recompute-min-nbor-dist does not apply to this model: its "
+                "tabulation domain is fixed analytically, so no neighbor "
+                "statistics are measured. Ignoring the flag."
+            )
+        else:
+            log.info(
+                "The model tabulates over an analytically bounded domain; "
+                "skipping the neighbor statistics."
+            )
+        # Carry a measured distance through untouched; leave it absent when
+        # the model has none, rather than inventing one.
+        min_nbor_dist = saved_min_nbor_dist
+        if min_nbor_dist is not None:
+            log.info(f"Minimal neighbor distance read from {source}: {min_nbor_dist:f}")
 
-    model.min_nbor_dist = min_nbor_dist
+    # ``enable_compression`` reads the value off the model, so a descriptor
+    # that ignores it still needs a number there. That placeholder stays local
+    # to the call: afterwards the model carries a measured distance or none.
+    model.min_nbor_dist = 0.0 if min_nbor_dist is None else 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)
@@ -110,6 +185,7 @@ def enable_compression(
         stride * 10,
         check_frequency,
     )
+    model.min_nbor_dist = min_nbor_dist
 
     # 4. Serialize the compressed model dict (includes tabulated data)
     compressed_model_dict = model.serialize()
@@ -130,11 +206,18 @@ def enable_compression(
     )
 
     model_def_script = model_dict.get("model_def_script")
+    # Only a measured distance is metadata; a model that never measured one
+    # leaves the field out rather than carrying a placeholder.
+    measured = {} if min_nbor_dist is None else {"min_nbor_dist": float(min_nbor_dist)}
     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},
+            {
+                "model": compressed_model_dict,
+                "model_def_script": model_def_script,
+                **measured,
+            },
             lower_kind="auto",
         )
     else:
@@ -145,7 +228,7 @@ def enable_compression(
             model_json_override={
                 "model": compressed_model_dict,
                 "model_def_script": model_def_script,
-                "min_nbor_dist": float(min_nbor_dist),
+                **measured,
             },
         )
     log.info("Compressed model saved to %s", output)

deepmd/pt_expt/entrypoints/main.py

@@ -970,6 +970,7 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None:
             extrapolate=FLAGS.extrapolate,
             check_frequency=FLAGS.frequency,
             training_script=FLAGS.training_script,
+            recompute_min_nbor_dist=FLAGS.recompute_min_nbor_dist,
         )
     else:
         raise RuntimeError(

deepmd/pt_expt/model/spin_model.py

@@ -48,6 +48,20 @@ def __getattr__(self, name: str) -> Any:
             return getattr(backbone, name)
         raise AttributeError(name)
 
+    @property
+    def min_nbor_dist(self) -> float | None:
+        """Minimal neighbor distance, stored on the backbone model.
+
+        ``__getattr__`` only delegates reads, so without this property an
+        assignment would land on the wrapper while ``get_min_nbor_dist`` and
+        ``enable_compression`` keep reading the backbone.
+        """
+        return self.backbone_model.min_nbor_dist
+
+    @min_nbor_dist.setter
+    def min_nbor_dist(self, value: float | None) -> None:
+        self.backbone_model.min_nbor_dist = value
+
     def forward_common_lower_exportable(
         self,
         extended_coord: torch.Tensor,

deepmd/pt_expt/utils/auto_batch_size.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+from deepmd.pt.utils.auto_batch_size import AutoBatchSize as AutoBatchSizeBase
+from deepmd.pt_expt.utils import (
+    env,
+)
+
+
+class AutoBatchSize(AutoBatchSizeBase):
+    """Auto batch size following the device pt_expt runs on.
+
+    ``env.DEVICE`` is CPU whenever ``DEVICE=cpu`` is set, even on a CUDA host.
+    Growing the batch there risks a host OOM, which the CUDA-OOM handler
+    cannot recover from, so the growth policy follows the selected device
+    rather than CUDA availability. The module is held rather than the value
+    so that the check reads the device in effect at call time.
+    """
+
+    def is_gpu_available(self) -> bool:
+        """Check if the selected device is a GPU.
+
+        Returns
+        -------
+        bool
+            True if pt_expt runs on a CUDA device
+        """
+        return env.DEVICE.type == "cuda"

deepmd/pt_expt/utils/neighbor_stat.py

@@ -10,6 +10,9 @@
 from deepmd.pt_expt.common import (
     torch_module,
 )
+from deepmd.pt_expt.utils.auto_batch_size import (
+    AutoBatchSize,
+)
 from deepmd.pt_expt.utils.env import (
     DEVICE,
     GLOBAL_PT_FLOAT_PRECISION,
@@ -28,6 +31,12 @@ class NeighborStatOP(NeighborStatOPDP):
 class NeighborStat(BaseNeighborStat):
     """Neighbor statistics using torch on DEVICE.
 
+    The statistics are evaluated one frame batch at a time. The intermediate
+    tensor is of shape ``[nframes, nloc, nall, 3]``, so processing a whole set
+    at once needs hundreds of GiB for a large set; :class:`AutoBatchSize`
+    keeps the batch within the available device memory, as the PyTorch,
+    Paddle, JAX and TensorFlow backends already do.
+
     Parameters
     ----------
     ntypes : int
@@ -46,6 +55,7 @@ def __init__(
     ) -> None:
         super().__init__(ntypes, rcut, mixed_type)
         self.op = NeighborStatOP(ntypes, rcut, mixed_type)
+        self.auto_batch_size = AutoBatchSize()
 
     def iterator(
         self, data: DeepmdDataSystem
@@ -65,7 +75,10 @@ def iterator(
             for jj in data.data_systems[ii].dirs:
                 data_set = data.data_systems[ii]
                 data_set_data = data_set._load_set(jj)
-                minrr2, max_nnei = self._execute(
+                minrr2, max_nnei = self.auto_batch_size.execute_all(
+                    self._execute,
+                    data_set_data["coord"].shape[0],
+                    data_set.get_natoms(),
                     data_set_data["coord"],
                     data_set_data["type"],
                     data_set_data["box"] if data_set.pbc else None,

source/tests/consistent/test_neighbor_stat.py

@@ -16,6 +16,7 @@
     INSTALLED_JAX,
     INSTALLED_PD,
     INSTALLED_PT,
+    INSTALLED_PT_EXPT,
     INSTALLED_TF,
 )
 
@@ -82,6 +83,10 @@ def test_neighbor_stat_tf(self) -> None:
     def test_neighbor_stat_pt(self) -> None:
         self.run_neighbor_stat("pytorch")
 
+    @unittest.skipUnless(INSTALLED_PT_EXPT, "PyTorch Exportable is not installed")
+    def test_neighbor_stat_pt_expt(self) -> None:
+        self.run_neighbor_stat("pt-expt")
+
     def test_neighbor_stat_dp(self) -> None:
         self.run_neighbor_stat("numpy")