From 15a0b342bd50d32c06cde55b9550f9d0799e53bd Mon Sep 17 00:00:00 2001 From: yckbz Date: Tue, 4 Aug 2026 12:33:57 +0800 Subject: [PATCH 01/15] fix(pt_expt): read min_nbor_dist from @variables dp convert-backend stores it there, so compress recomputed it from the training data on every run. Add --recompute-min-nbor-dist to force a recompute, and log where the value comes from. --- deepmd/main.py | 7 ++ deepmd/pt_expt/entrypoints/compress.py | 74 +++++++++++++++---- deepmd/pt_expt/entrypoints/main.py | 1 + .../pt_expt/test_compress_min_nbor_dist.py | 70 ++++++++++++++++++ 4 files changed, 139 insertions(+), 13 deletions(-) create mode 100644 source/tests/pt_expt/test_compress_min_nbor_dist.py diff --git a/deepmd/main.py b/deepmd/main.py index 1206b2fa0e..92377ed5e6 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -668,6 +668,13 @@ 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", + ) parser_compress.add_argument( "--head", "--model-branch", diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 3d9bb6506a..92dfa947a9 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/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,6 +14,33 @@ 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 enable_compression( input_file: str, output: str, @@ -18,6 +48,7 @@ def enable_compression( 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 +67,9 @@ 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. """ from deepmd.pt_expt.model.model import ( BaseModel, @@ -46,20 +80,34 @@ def enable_compression( model = BaseModel.deserialize(model_dict["model"]) # 2. Get or compute min_nbor_dist - 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: - 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)." + if recompute_min_nbor_dist: + min_nbor_dist, source = None, "" + else: + min_nbor_dist, source = _read_saved_min_nbor_dist(model, model_dict) + if min_nbor_dist is not None: + log.info(f"Minimal neighbor distance read from {source}: {min_nbor_dist:f}") + else: + if recompute_min_nbor_dist: + log.info( + "Recompute the minimal neighbor distance from the training data, " + "ignoring the one saved in the model." + ) + if training_script is None: + raise ValueError( + "Recomputing the minimal neighbor distance requires the " + "training script and data (via -t,--training-script)." + ) + 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)." + ) from deepmd.common import ( j_loader, ) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 66ddd1de79..5e9f6addf3 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -928,6 +928,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( diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py new file mode 100644 index 0000000000..7f090d2fe3 --- /dev/null +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for reading the stored minimal neighbor distance in pt_expt compress.""" + +from typing import ( + Any, +) + +import pytest + +from deepmd.main import ( + parse_args, +) +from deepmd.pt_expt.entrypoints.compress import ( + _read_saved_min_nbor_dist, +) + + +class _FakeModel: + """Stand-in exposing only the getter that the reader calls.""" + + def __init__(self, min_nbor_dist: float | None) -> None: + self._min_nbor_dist = min_nbor_dist + + def get_min_nbor_dist(self) -> float | None: + return self._min_nbor_dist + + +@pytest.mark.parametrize( + ("model_min_nbor_dist", "model_dict", "expected"), + [ + # the model buffer wins over both metadata locations + ( + 1.5, + {"min_nbor_dist": 9.0, "@variables": {"min_nbor_dist": 8.0}}, + (1.5, "the model"), + ), + # the top-level key, written by pt_expt's own compress + ( + None, + {"min_nbor_dist": 2.5, "@variables": {"min_nbor_dist": 8.0}}, + (2.5, "the model file"), + ), + # "@variables", the cross-backend location that dp convert-backend writes + ( + None, + {"@variables": {"min_nbor_dist": 3.5}}, + (3.5, "the model file (@variables)"), + ), + # nothing stored anywhere + (None, {}, (None, "")), + (None, {"@variables": {}}, (None, "")), + (None, {"@variables": None}, (None, "")), + ], +) +def test_read_saved_min_nbor_dist( + model_min_nbor_dist: float | None, + model_dict: dict[str, Any], + expected: tuple[float | None, str], +) -> None: + """Every storage location is honored, in order of precedence.""" + assert _read_saved_min_nbor_dist(_FakeModel(model_min_nbor_dist), model_dict) == ( + expected + ) + + +def test_recompute_min_nbor_dist_flag_defaults_to_false() -> None: + """The compress parser exposes the flag and leaves it off by default.""" + args = ["--pt-expt", "compress", "-i", "in.pt2", "-o", "out.pt2"] + assert parse_args(args).recompute_min_nbor_dist is False + assert parse_args([*args, "--recompute-min-nbor-dist"]).recompute_min_nbor_dist From 110e60b23adba2b4eeeefb6e111e94c72170d4fc Mon Sep 17 00:00:00 2001 From: yckbz Date: Tue, 4 Aug 2026 12:33:57 +0800 Subject: [PATCH 02/15] fix(pt_expt): batch the neighbor statistics Sending a whole set to the device at once needs hundreds of GiB. Use AutoBatchSize, as the pt, pd, jax and tf backends already do. --- deepmd/pt_expt/utils/neighbor_stat.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/deepmd/pt_expt/utils/neighbor_stat.py b/deepmd/pt_expt/utils/neighbor_stat.py index 850bf86def..36f4788cc0 100644 --- a/deepmd/pt_expt/utils/neighbor_stat.py +++ b/deepmd/pt_expt/utils/neighbor_stat.py @@ -7,6 +7,9 @@ import torch from deepmd.dpmodel.utils.neighbor_stat import NeighborStatOP as NeighborStatOPDP +from deepmd.pt.utils.auto_batch_size import ( + AutoBatchSize, +) from deepmd.pt_expt.common import ( torch_module, ) @@ -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, From 6681b324a011e2e2552bc20fc8702cb048e7e377 Mon Sep 17 00:00:00 2001 From: Ce Liu Date: Tue, 4 Aug 2026 12:44:43 +0800 Subject: [PATCH 03/15] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Ce Liu --- deepmd/pt_expt/entrypoints/compress.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 92dfa947a9..f9086d17e6 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -14,7 +14,9 @@ log = logging.getLogger(__name__) -def _read_saved_min_nbor_dist(model: Any, model_dict: dict) -> tuple[float | None, str]: +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 From 12e34ad6c3dc7179052518a72925d751028133c0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 04:45:52 +0000 Subject: [PATCH 04/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt_expt/entrypoints/compress.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index f9086d17e6..92dfa947a9 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -14,9 +14,7 @@ log = logging.getLogger(__name__) -def _read_saved_min_nbor_dist( - model: Any, model_dict: dict -) -> tuple[float | None, str]: +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 From 0c88b5013506c5da0a80b676b08b6d95a587d33b Mon Sep 17 00:00:00 2001 From: yckbz Date: Tue, 11 Aug 2026 10:35:34 +0800 Subject: [PATCH 05/15] fix(pt_expt): follow the selected device when auto-batching DEVICE is CPU whenever DEVICE=cpu is set, even on a CUDA host, where the CUDA-OOM handler cannot recover from a host OOM. Base the growth policy on the selected device instead of CUDA availability. --- deepmd/pt_expt/utils/auto_batch_size.py | 28 +++++++++++++++++++ deepmd/pt_expt/utils/neighbor_stat.py | 6 ++-- .../pt_expt/test_compress_min_nbor_dist.py | 16 +++++++++++ 3 files changed, 47 insertions(+), 3 deletions(-) create mode 100644 deepmd/pt_expt/utils/auto_batch_size.py diff --git a/deepmd/pt_expt/utils/auto_batch_size.py b/deepmd/pt_expt/utils/auto_batch_size.py new file mode 100644 index 0000000000..c13ffcd377 --- /dev/null +++ b/deepmd/pt_expt/utils/auto_batch_size.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from deepmd.pt.utils.auto_batch_size import ( + AutoBatchSize as AutoBatchSizeBase, +) +from deepmd.pt_expt.utils.env import ( + DEVICE, +) + + +class AutoBatchSize(AutoBatchSizeBase): + """Auto batch size following the device pt_expt runs on. + + ``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. + """ + + 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 DEVICE.type == "cuda" diff --git a/deepmd/pt_expt/utils/neighbor_stat.py b/deepmd/pt_expt/utils/neighbor_stat.py index 36f4788cc0..2b09ae6a02 100644 --- a/deepmd/pt_expt/utils/neighbor_stat.py +++ b/deepmd/pt_expt/utils/neighbor_stat.py @@ -7,12 +7,12 @@ import torch from deepmd.dpmodel.utils.neighbor_stat import NeighborStatOP as NeighborStatOPDP -from deepmd.pt.utils.auto_batch_size import ( - AutoBatchSize, -) 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, diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py index 7f090d2fe3..15a3a85e4f 100644 --- a/source/tests/pt_expt/test_compress_min_nbor_dist.py +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -68,3 +68,19 @@ def test_recompute_min_nbor_dist_flag_defaults_to_false() -> None: args = ["--pt-expt", "compress", "-i", "in.pt2", "-o", "out.pt2"] assert parse_args(args).recompute_min_nbor_dist is False assert parse_args([*args, "--recompute-min-nbor-dist"]).recompute_min_nbor_dist + + +def test_auto_batch_size_follows_the_selected_device() -> None: + """Growth is allowed only when pt_expt actually runs on a GPU. + + ``DEVICE`` is CPU whenever ``DEVICE=cpu`` is set, even on a CUDA host, + where an expanding batch would risk an unrecoverable host OOM. + """ + from deepmd.pt_expt.utils.auto_batch_size import ( + AutoBatchSize, + ) + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + assert AutoBatchSize().is_gpu_available() == (DEVICE.type == "cuda") From d850ca9f430588d2192ac8122009518c90814b04 Mon Sep 17 00:00:00 2001 From: yckbz Date: Tue, 11 Aug 2026 10:36:17 +0800 Subject: [PATCH 06/15] fix(pt_expt): keep min_nbor_dist in the graph-lower output model.serialize() does not carry the runtime buffer, so the graph branch dropped the value it had just recovered and the next run fell back to -t. --- deepmd/pt_expt/entrypoints/compress.py | 6 ++- .../pt_expt/model/test_model_compression.py | 29 +++++++++++++ .../pt_expt/test_compress_min_nbor_dist.py | 41 +++++++++++++++++++ 3 files changed, 75 insertions(+), 1 deletion(-) diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 92dfa947a9..82ffbc3f88 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -174,7 +174,11 @@ def enable_compression( 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, + "min_nbor_dist": float(min_nbor_dist), + }, lower_kind="auto", ) else: diff --git a/source/tests/pt_expt/model/test_model_compression.py b/source/tests/pt_expt/model/test_model_compression.py index 95a66b528d..2c84e54049 100644 --- a/source/tests/pt_expt/model/test_model_compression.py +++ b/source/tests/pt_expt/model/test_model_compression.py @@ -187,6 +187,35 @@ def test_min_nbor_dist_roundtrip(self) -> None: finally: os.unlink(frozen_path) + def test_compress_min_nbor_dist_from_variables(self) -> None: + """Test that compress recovers min_nbor_dist from @variables and keeps it. + + ``dp convert-backend`` stores the value under ``@variables``. Compress + must read it from there — without a training script it would otherwise + raise — and carry it into the compressed archive. + """ + from deepmd.pt_expt.entrypoints.compress import ( + enable_compression, + ) + + md = self._make_model() + md.eval() + + model_data = {"model": md.serialize(), "@variables": {"min_nbor_dist": 0.5}} + with tempfile.NamedTemporaryFile(suffix=".pte", delete=False) as f: + frozen_path = f.name + with tempfile.NamedTemporaryFile(suffix=".pte", delete=False) as f: + compressed_path = f.name + try: + deserialize_to_file(frozen_path, model_data) + enable_compression(input_file=frozen_path, output=compressed_path) + compressed_data = serialize_from_file(compressed_path) + self.assertAlmostEqual(compressed_data["min_nbor_dist"], 0.5) + finally: + os.unlink(frozen_path) + if os.path.exists(compressed_path): + os.unlink(compressed_path) + def test_compress_state_serialized(self) -> None: """Test that compression state persists through serialize/deserialize. diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py index 15a3a85e4f..dbbe50fcff 100644 --- a/source/tests/pt_expt/test_compress_min_nbor_dist.py +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -84,3 +84,44 @@ def test_auto_batch_size_follows_the_selected_device() -> None: ) assert AutoBatchSize().is_gpu_available() == (DEVICE.type == "cuda") + + +def test_graph_lower_output_keeps_min_nbor_dist(monkeypatch: Any) -> None: + """The graph-lower branch carries the value into the exported archive. + + ``model.serialize()`` does not include the runtime buffer, so without this + the compressed artifact loses a value that was just recovered. + """ + import deepmd.pt_expt.entrypoints.compress as compress_mod + import deepmd.pt_expt.model.graph_lower as graph_lower_mod + import deepmd.pt_expt.model.model as model_mod + import deepmd.pt_expt.utils.tabulate_ops as tabulate_ops_mod + + class _Model(_FakeModel): + def enable_compression(self, *args: Any, **kwargs: Any) -> None: + pass + + def serialize(self) -> dict: + return {"type": "compressed"} + + captured: dict[str, Any] = {} + + monkeypatch.setattr( + compress_mod, + "serialize_from_file", + lambda _: {"model": {}, "@variables": {"min_nbor_dist": 1.25}}, + ) + monkeypatch.setattr( + model_mod.BaseModel, "deserialize", staticmethod(lambda _: _Model(None)) + ) + monkeypatch.setattr(tabulate_ops_mod, "ensure_fake_registered", lambda: None) + monkeypatch.setattr(graph_lower_mod, "model_uses_graph_lower", lambda _: True) + monkeypatch.setattr( + compress_mod, + "deserialize_to_file", + lambda output, data, **kwargs: captured.update(data=data), + ) + + compress_mod.enable_compression(input_file="in.pt2", output="out.pt2") + + assert captured["data"]["min_nbor_dist"] == 1.25 From f1423c1c952fcec38c7ab6092f951bb25c9b0d0f Mon Sep 17 00:00:00 2001 From: yckbz Date: Tue, 11 Aug 2026 10:36:17 +0800 Subject: [PATCH 07/15] fix(pt_expt): assign min_nbor_dist through the spin backbone SpinModel.__getattr__ only delegates reads, so the assignment landed on the wrapper while get_min_nbor_dist and enable_compression kept reading the backbone. --- deepmd/pt_expt/model/spin_model.py | 14 ++++++++++ .../pt_expt/test_compress_min_nbor_dist.py | 26 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/deepmd/pt_expt/model/spin_model.py b/deepmd/pt_expt/model/spin_model.py index 83cc8ecba4..b4e9a8bcd0 100644 --- a/deepmd/pt_expt/model/spin_model.py +++ b/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, diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py index dbbe50fcff..273d3c0dd0 100644 --- a/source/tests/pt_expt/test_compress_min_nbor_dist.py +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -125,3 +125,29 @@ def serialize(self) -> dict: compress_mod.enable_compression(input_file="in.pt2", output="out.pt2") assert captured["data"]["min_nbor_dist"] == 1.25 + + +def test_spin_model_min_nbor_dist_reaches_the_backbone() -> None: + """Assigning on the wrapper must reach the model that gets compressed. + + ``SpinModel.__getattr__`` only delegates reads, while both + ``get_min_nbor_dist`` and ``enable_compression`` run on the backbone. + """ + from deepmd.pt_expt.model.spin_model import ( + SpinModel, + ) + + class _Backbone: + min_nbor_dist: float | None = None + + class _Wrapper: + # the property under test, on a stand-in that needs no torch module + min_nbor_dist = SpinModel.min_nbor_dist + + def __init__(self) -> None: + self.backbone_model = _Backbone() + + wrapper = _Wrapper() + wrapper.min_nbor_dist = 1.25 + assert wrapper.backbone_model.min_nbor_dist == 1.25 + assert wrapper.min_nbor_dist == 1.25 From cbb5117541b20e532609eef13bb5f9b39c77d0e7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 11 Aug 2026 02:57:34 +0000 Subject: [PATCH 08/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt_expt/utils/auto_batch_size.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/deepmd/pt_expt/utils/auto_batch_size.py b/deepmd/pt_expt/utils/auto_batch_size.py index c13ffcd377..ed23e89550 100644 --- a/deepmd/pt_expt/utils/auto_batch_size.py +++ b/deepmd/pt_expt/utils/auto_batch_size.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from deepmd.pt.utils.auto_batch_size import ( - AutoBatchSize as AutoBatchSizeBase, -) +from deepmd.pt.utils.auto_batch_size import AutoBatchSize as AutoBatchSizeBase from deepmd.pt_expt.utils.env import ( DEVICE, ) From 913f68a7ff06789951599b5ef6275954e757246a Mon Sep 17 00:00:00 2001 From: yckbz Date: Sat, 29 Aug 2026 09:23:51 +0800 Subject: [PATCH 09/15] fix(pt_expt): read DEVICE at call time in the auto batch size Importing the value froze the device at import time, which defeats the point of a policy meant to follow the selected device. --- deepmd/pt_expt/utils/auto_batch_size.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/deepmd/pt_expt/utils/auto_batch_size.py b/deepmd/pt_expt/utils/auto_batch_size.py index ed23e89550..f1b540d302 100644 --- a/deepmd/pt_expt/utils/auto_batch_size.py +++ b/deepmd/pt_expt/utils/auto_batch_size.py @@ -1,18 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from deepmd.pt.utils.auto_batch_size import AutoBatchSize as AutoBatchSizeBase -from deepmd.pt_expt.utils.env import ( - DEVICE, +from deepmd.pt_expt.utils import ( + env, ) class AutoBatchSize(AutoBatchSizeBase): """Auto batch size following the device pt_expt runs on. - ``DEVICE`` is CPU whenever ``DEVICE=cpu`` is set, even on a CUDA host. + ``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. + 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: @@ -23,4 +24,4 @@ def is_gpu_available(self) -> bool: bool True if pt_expt runs on a CUDA device """ - return DEVICE.type == "cuda" + return env.DEVICE.type == "cuda" From 9a914becbd4b29c050fdd6ed481176165442786d Mon Sep 17 00:00:00 2001 From: yckbz Date: Sat, 29 Aug 2026 09:23:51 +0800 Subject: [PATCH 10/15] test(pt_expt): import the compress module one way It was reached both by `import` and by `import from`. --- source/tests/pt_expt/test_compress_min_nbor_dist.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py index 273d3c0dd0..c931b894fa 100644 --- a/source/tests/pt_expt/test_compress_min_nbor_dist.py +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -10,8 +10,8 @@ from deepmd.main import ( parse_args, ) -from deepmd.pt_expt.entrypoints.compress import ( - _read_saved_min_nbor_dist, +from deepmd.pt_expt.entrypoints import ( + compress as compress_mod, ) @@ -58,9 +58,10 @@ def test_read_saved_min_nbor_dist( expected: tuple[float | None, str], ) -> None: """Every storage location is honored, in order of precedence.""" - assert _read_saved_min_nbor_dist(_FakeModel(model_min_nbor_dist), model_dict) == ( - expected + got = compress_mod._read_saved_min_nbor_dist( + _FakeModel(model_min_nbor_dist), model_dict ) + assert got == expected def test_recompute_min_nbor_dist_flag_defaults_to_false() -> None: @@ -92,7 +93,6 @@ def test_graph_lower_output_keeps_min_nbor_dist(monkeypatch: Any) -> None: ``model.serialize()`` does not include the runtime buffer, so without this the compressed artifact loses a value that was just recovered. """ - import deepmd.pt_expt.entrypoints.compress as compress_mod import deepmd.pt_expt.model.graph_lower as graph_lower_mod import deepmd.pt_expt.model.model as model_mod import deepmd.pt_expt.utils.tabulate_ops as tabulate_ops_mod From 2c69cffde827088b4c8e4f5845cb84e5d5cf2f53 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:25:03 +0000 Subject: [PATCH 11/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/pt_expt/test_compress_min_nbor_dist.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py index c931b894fa..0be625366e 100644 --- a/source/tests/pt_expt/test_compress_min_nbor_dist.py +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -10,9 +10,7 @@ from deepmd.main import ( parse_args, ) -from deepmd.pt_expt.entrypoints import ( - compress as compress_mod, -) +from deepmd.pt_expt.entrypoints import compress as compress_mod class _FakeModel: From 3b3f49303cec40269d63b081e4b8e5a00b1787e5 Mon Sep 17 00:00:00 2001 From: yckbz Date: Thu, 15 Oct 2026 19:00:37 +0800 Subject: [PATCH 12/15] fix(pt_expt): keep the compression placeholder out of the metadata --recompute-min-nbor-dist dropped the saved distance before checking whether compression consumes it, so a model that skips the statistics wrote the 0.0 placeholder over a valid one. Decide on that first, and record only a measured distance. --- deepmd/main.py | 4 +- deepmd/pt_expt/entrypoints/compress.py | 119 ++++++++----- .../pt_expt/test_compress_min_nbor_dist.py | 168 +++++++++++++++++- 3 files changed, 244 insertions(+), 47 deletions(-) diff --git a/deepmd/main.py b/deepmd/main.py index 92377ed5e6..d0b45e3260 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -673,7 +673,9 @@ def main_parser() -> argparse.ArgumentParser: 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", + "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", diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index d028aaefce..016489aa00 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -41,6 +41,40 @@ def _read_saved_min_nbor_dist(model: Any, model_dict: dict) -> tuple[float | Non 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, @@ -69,7 +103,10 @@ def enable_compression( 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. + 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, @@ -79,22 +116,11 @@ 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. - if recompute_min_nbor_dist: - min_nbor_dist, source = None, "" - else: - min_nbor_dist, source = _read_saved_min_nbor_dist(model, model_dict) - if min_nbor_dist is not None: - log.info(f"Minimal neighbor distance read from {source}: {min_nbor_dist:f}") - elif 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 - else: + # 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, " @@ -105,6 +131,10 @@ def enable_compression( "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, " @@ -116,32 +146,29 @@ def enable_compression( "so the training script and data must be provided " "(via -t,--training-script)." ) - 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) + 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) @@ -158,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() @@ -178,6 +206,9 @@ 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( @@ -185,7 +216,7 @@ def enable_compression( { "model": compressed_model_dict, "model_def_script": model_def_script, - "min_nbor_dist": float(min_nbor_dist), + **measured, }, lower_kind="auto", ) @@ -197,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) diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py index 0be625366e..2f69a05e20 100644 --- a/source/tests/pt_expt/test_compress_min_nbor_dist.py +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Tests for reading the stored minimal neighbor distance in pt_expt compress.""" +import logging from typing import ( Any, ) @@ -14,14 +15,23 @@ class _FakeModel: - """Stand-in exposing only the getter that the reader calls.""" + """Stand-in exposing only what the compression entry point calls.""" - def __init__(self, min_nbor_dist: float | None) -> None: + def __init__( + self, + min_nbor_dist: float | None, + *, + needs_min_nbor_dist: bool = True, + ) -> None: self._min_nbor_dist = min_nbor_dist + self._needs_min_nbor_dist = needs_min_nbor_dist def get_min_nbor_dist(self) -> float | None: return self._min_nbor_dist + def compression_needs_min_nbor_dist(self) -> bool: + return self._needs_min_nbor_dist + @pytest.mark.parametrize( ("model_min_nbor_dist", "model_dict", "expected"), @@ -149,3 +159,157 @@ def __init__(self) -> None: wrapper.min_nbor_dist = 1.25 assert wrapper.backbone_model.min_nbor_dist == 1.25 assert wrapper.min_nbor_dist == 1.25 + + +def _run_compression( + monkeypatch: Any, + *, + saved: float | None, + needs: bool, + flag: bool, + training_script: str | None = "train.json", +) -> tuple[dict, list, Any]: + """Compress a stand-in model and report what reached the archive. + + Returns + ------- + dict + The data handed to ``deserialize_to_file``. + list + One entry per training-data scan, so an empty list proves none ran. + Any + The model, to inspect what the run left on it. + """ + import deepmd.pt_expt.model.graph_lower as graph_lower_mod + import deepmd.pt_expt.model.model as model_mod + import deepmd.pt_expt.utils.tabulate_ops as tabulate_ops_mod + + class _Model(_FakeModel): + def enable_compression(self, *args: Any, **kwargs: Any) -> None: + pass + + def serialize(self) -> dict: + return {"type": "compressed"} + + model = _Model(None, needs_min_nbor_dist=needs) + model_dict: dict[str, Any] = {"model": {}} + if saved is not None: + model_dict["@variables"] = {"min_nbor_dist": saved} + + captured: dict[str, Any] = {} + scans: list[str] = [] + + monkeypatch.setattr(compress_mod, "serialize_from_file", lambda _: model_dict) + monkeypatch.setattr( + model_mod.BaseModel, "deserialize", staticmethod(lambda _: model) + ) + monkeypatch.setattr(tabulate_ops_mod, "ensure_fake_registered", lambda: None) + monkeypatch.setattr(graph_lower_mod, "model_uses_graph_lower", lambda _: True) + monkeypatch.setattr( + compress_mod, + "deserialize_to_file", + lambda output, data, **kwargs: captured.update(data=data), + ) + monkeypatch.setattr( + compress_mod, + "_measure_min_nbor_dist", + lambda script: (scans.append(script), 9.99)[1], + ) + + compress_mod.enable_compression( + input_file="in.pt2", + output="out.pt2", + training_script=training_script, + recompute_min_nbor_dist=flag, + ) + return captured["data"], scans, model + + +@pytest.mark.parametrize("flag", [False, True]) +def test_unused_statistic_keeps_a_saved_distance(monkeypatch: Any, flag: bool) -> None: + """A model that ignores the statistic must not lose a measured one. + + The recompute flag is inapplicable here, so it may not replace a real + distance with the placeholder the compression ABI is given. + """ + data, scans, model = _run_compression( + monkeypatch, saved=0.75, needs=False, flag=flag + ) + + assert data["min_nbor_dist"] == 0.75 + assert scans == [] + assert model.min_nbor_dist == 0.75 + + +@pytest.mark.parametrize("flag", [False, True]) +def test_unused_statistic_leaves_an_absent_distance_absent( + monkeypatch: Any, flag: bool +) -> None: + """Never measured means absent, not zero. + + ``enable_compression`` still needs a number, but that placeholder stays + local to the call instead of being recorded as a measurement. + """ + data, scans, model = _run_compression( + monkeypatch, saved=None, needs=False, flag=flag + ) + + assert "min_nbor_dist" not in data + assert scans == [] + assert model.min_nbor_dist is None + + +def test_unused_statistic_warns_the_flag_is_inapplicable( + monkeypatch: Any, caplog: Any +) -> None: + """The flag is ignored loudly, not silently.""" + with caplog.at_level(logging.WARNING): + _run_compression(monkeypatch, saved=0.75, needs=False, flag=True) + + assert "--recompute-min-nbor-dist does not apply" in caplog.text + + +@pytest.mark.parametrize( + ("saved", "flag", "expected", "scanned"), + [ + # a saved value is reused, so no pass over the training data + (1.25, False, 1.25, False), + # an explicit request measures again and wins over the saved value + (1.25, True, 9.99, True), + # nothing saved leaves measuring as the only option + (None, False, 9.99, True), + ], +) +def test_consumed_statistic_paths( + monkeypatch: Any, + saved: float | None, + flag: bool, + expected: float, + scanned: bool, +) -> None: + """A model whose tabulation consumes the statistic is unaffected.""" + data, scans, _ = _run_compression(monkeypatch, saved=saved, needs=True, flag=flag) + + assert data["min_nbor_dist"] == expected + assert bool(scans) is scanned + + +@pytest.mark.parametrize( + ("flag", "message"), + [ + (True, "Recomputing the minimal neighbor distance requires"), + (False, "does not have a minimum neighbor distance"), + ], +) +def test_measuring_requires_a_training_script( + monkeypatch: Any, flag: bool, message: str +) -> None: + """Both routes into measuring keep their training-script validation.""" + with pytest.raises(ValueError, match=message): + _run_compression( + monkeypatch, + saved=None if not flag else 1.25, + needs=True, + flag=flag, + training_script=None, + ) From 88bf6a8eefb7839d43eebc667c947a2425194867 Mon Sep 17 00:00:00 2001 From: yckbz Date: Thu, 15 Oct 2026 19:13:20 +0800 Subject: [PATCH 13/15] test(pt_expt): cover the batched neighbor statistics Nothing in source/tests imported the pt_expt NeighborStat and the consistent suite skipped that backend, so the batching had no test. Pin the split itself, since comparing results alone passes unbatched too. Also run the consistent suite against pt-expt, and recompute with nothing saved. --- source/tests/consistent/test_neighbor_stat.py | 4 + .../pt_expt/test_compress_min_nbor_dist.py | 2 + .../tests/pt_expt/utils/test_neighbor_stat.py | 138 ++++++++++++++++++ 3 files changed, 144 insertions(+) create mode 100644 source/tests/pt_expt/utils/test_neighbor_stat.py diff --git a/source/tests/consistent/test_neighbor_stat.py b/source/tests/consistent/test_neighbor_stat.py index 9c9f97045b..35ff97ae37 100644 --- a/source/tests/consistent/test_neighbor_stat.py +++ b/source/tests/consistent/test_neighbor_stat.py @@ -82,6 +82,10 @@ def test_neighbor_stat_tf(self) -> None: def test_neighbor_stat_pt(self) -> None: self.run_neighbor_stat("pytorch") + @unittest.skipUnless(INSTALLED_PT, "pytorch 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") diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py index 2f69a05e20..b12d93b65d 100644 --- a/source/tests/pt_expt/test_compress_min_nbor_dist.py +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -278,6 +278,8 @@ def test_unused_statistic_warns_the_flag_is_inapplicable( (1.25, True, 9.99, True), # nothing saved leaves measuring as the only option (None, False, 9.99, True), + # an explicit request with nothing saved measures just the same + (None, True, 9.99, True), ], ) def test_consumed_statistic_paths( diff --git a/source/tests/pt_expt/utils/test_neighbor_stat.py b/source/tests/pt_expt/utils/test_neighbor_stat.py new file mode 100644 index 0000000000..6101f2c6f6 --- /dev/null +++ b/source/tests/pt_expt/utils/test_neighbor_stat.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the batched neighbor statistics of the pt_expt backend.""" + +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.pt_expt.utils.neighbor_stat import ( + NeighborStat, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) + +# A 3x3x3 grid of one atom type, spaced 1.0 apart in a 3.0 box, so the +# shortest neighbor distance is exactly 1.0 and every atom has the same +# neighborhood regardless of which frames land in a batch. +NATOMS = 27 +BOX_LENGTH = 3.0 +NFRAMES = 8 +SET_SIZE = 4 +RCUT = 1.001 + + +def _write_system(path: Path) -> None: + """Write a multi-frame, multi-set system in the DeePMD npy layout.""" + grid = np.mgrid[0:2:3j, 0:2:3j, 0:2:3j] + positions = np.vstack([axis.ravel() for axis in grid]).T + coord = np.repeat(positions[np.newaxis, :, :], NFRAMES, axis=0).reshape( + NFRAMES, NATOMS * 3 + ) + box = np.tile((np.eye(3) * BOX_LENGTH).ravel(), (NFRAMES, 1)) + + path.mkdir(parents=True, exist_ok=True) + (path / "type.raw").write_text("\n".join(["0"] * NATOMS) + "\n") + (path / "type_map.raw").write_text("TYPE\n") + for index, start in enumerate(range(0, NFRAMES, SET_SIZE)): + subset = path / f"set.{index:03d}" + subset.mkdir() + frames = slice(start, start + SET_SIZE) + np.save(subset / "coord.npy", coord[frames]) + np.save(subset / "box.npy", box[frames]) + np.save(subset / "energy.npy", np.zeros(SET_SIZE)) + np.save(subset / "force.npy", np.zeros((SET_SIZE, NATOMS * 3))) + + +def _stat_with_batch_size( + monkeypatch: Any, system: str, batch_atoms: int +) -> tuple[tuple[float, np.ndarray], list[int]]: + """Run the statistics at a pinned batch size. + + ``DP_INFER_BATCH_SIZE`` is read when :class:`AutoBatchSize` is built, so + it has to be set before the ``NeighborStat`` is constructed. + + Returns + ------- + tuple + The ``(min_nbor_dist, max_nbor_size)`` that ``get_stat`` produced. + list + The number of frames handed to each ``_execute`` call, one entry per + call, so the caller can tell how the work was split. + """ + frames_per_call: list[int] = [] + # A context keeps the patched ``_execute`` from leaking into a later run, + # which would otherwise chain the spies and mix their records. + with monkeypatch.context() as patch: + patch.setenv("DP_INFER_BATCH_SIZE", str(batch_atoms)) + original = NeighborStat._execute + + def spy( + self: NeighborStat, + coord: np.ndarray, + atype: np.ndarray, + cell: np.ndarray | None, + ) -> tuple[np.ndarray, np.ndarray]: + frames_per_call.append(coord.shape[0]) + return original(self, coord, atype, cell) + + patch.setattr(NeighborStat, "_execute", spy) + + data = DeepmdDataSystem( + systems=[system], + batch_size=1, + test_size=1, + rcut=RCUT, + type_map=["TYPE"], + ) + data.get_batch() + stat = NeighborStat(data.get_ntypes(), RCUT, mixed_type=False).get_stat(data) + return stat, frames_per_call + + +def test_batching_splits_the_work_and_keeps_the_result( + monkeypatch: Any, tmp_path: Path +) -> None: + """Batching must change how the work is split, not what it produces. + + Sending a whole set at once allocates an ``[nframes, nloc, nall, 3]`` + tensor, which is why the statistics are batched. Comparing the results + alone would pass without any batching, so the split itself is asserted. + """ + system = tmp_path / "system" + _write_system(system) + + whole, whole_calls = _stat_with_batch_size( + monkeypatch, str(system), NATOMS * NFRAMES + ) + batched, batched_calls = _stat_with_batch_size(monkeypatch, str(system), NATOMS) + + # one call per set when everything fits, one call per frame when it does not + assert whole_calls == [SET_SIZE] * (NFRAMES // SET_SIZE) + assert batched_calls == [1] * NFRAMES + assert len(batched_calls) > len(whole_calls) + + assert batched[0] == whole[0] + np.testing.assert_array_equal(batched[1], whole[1]) + + +def test_batched_statistics_are_correct(monkeypatch: Any, tmp_path: Path) -> None: + """The reduction over batches reproduces the known geometry. + + On a unit grid the shortest distance is 1.0, and within ``rcut`` each + atom sees its six axis-aligned images. + """ + system = tmp_path / "system" + _write_system(system) + + (min_nbor_dist, max_nbor_size), _ = _stat_with_batch_size( + monkeypatch, str(system), NATOMS + ) + + assert min_nbor_dist == 1.0 + np.testing.assert_array_equal(max_nbor_size, [6]) From fc01e96b6184006263a7dd4ca566b6eabd3750e2 Mon Sep 17 00:00:00 2001 From: yckbz Date: Sun, 18 Oct 2026 13:46:28 +0800 Subject: [PATCH 14/15] test(pt_expt): make the two new tests tell the fixes apart The batching frames were identical, so a reduction keeping one batch would still match; the extrema now come from different sets. The device check was an identity that both policies satisfy; pin the CUDA host with cpu selected, the one combination where they differ. --- .../pt_expt/test_compress_min_nbor_dist.py | 22 +++ .../tests/pt_expt/utils/test_neighbor_stat.py | 133 +++++++++++++----- 2 files changed, 122 insertions(+), 33 deletions(-) diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py index b12d93b65d..b7a9dd4e74 100644 --- a/source/tests/pt_expt/test_compress_min_nbor_dist.py +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -95,6 +95,28 @@ def test_auto_batch_size_follows_the_selected_device() -> None: assert AutoBatchSize().is_gpu_available() == (DEVICE.type == "cuda") +def test_auto_batch_size_ignores_cuda_when_cpu_is_selected(monkeypatch: Any) -> None: + """A CUDA host with ``DEVICE=cpu`` must still pin the CPU policy. + + This is the combination the policy exists for, and the only one where + following ``torch.cuda.is_available()`` instead of the selected device + gives a different answer. + """ + import torch + + from deepmd.pt_expt.utils import ( + env, + ) + from deepmd.pt_expt.utils.auto_batch_size import ( + AutoBatchSize, + ) + + monkeypatch.setattr(env, "DEVICE", torch.device("cpu")) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + assert AutoBatchSize(silent=True).is_gpu_available() is False + + def test_graph_lower_output_keeps_min_nbor_dist(monkeypatch: Any) -> None: """The graph-lower branch carries the value into the exported archive. diff --git a/source/tests/pt_expt/utils/test_neighbor_stat.py b/source/tests/pt_expt/utils/test_neighbor_stat.py index 6101f2c6f6..2a0af3d210 100644 --- a/source/tests/pt_expt/utils/test_neighbor_stat.py +++ b/source/tests/pt_expt/utils/test_neighbor_stat.py @@ -17,40 +17,84 @@ DeepmdDataSystem, ) -# A 3x3x3 grid of one atom type, spaced 1.0 apart in a 3.0 box, so the -# shortest neighbor distance is exactly 1.0 and every atom has the same -# neighborhood regardless of which frames land in a batch. NATOMS = 27 -BOX_LENGTH = 3.0 NFRAMES = 8 SET_SIZE = 4 -RCUT = 1.001 +# A 3x3x3 grid spaced 1.0 apart in a 3.0 box: the shortest distance is +# exactly 1.0 and, just past it, every atom sees its six axis-aligned +# images. Used for the analytical check. +UNIFORM_SPACING = 1.0 +UNIFORM_RCUT = 1.001 -def _write_system(path: Path) -> None: - """Write a multi-frame, multi-set system in the DeePMD npy layout.""" - grid = np.mgrid[0:2:3j, 0:2:3j, 0:2:3j] - positions = np.vstack([axis.ravel() for axis in grid]).T - coord = np.repeat(positions[np.newaxis, :, :], NFRAMES, axis=0).reshape( - NFRAMES, NATOMS * 3 - ) - box = np.tile((np.eye(3) * BOX_LENGTH).ravel(), (NFRAMES, 1)) +# For the batching comparison the cutoff also admits the second shell +# (sqrt(2) * spacing), which is what makes the neighbor count depend on +# the spacing of each individual frame. +MIXED_RCUT = 1.5 +SPARSE_SPACING = 1.2 +DENSE_SPACING = 1.0 +CLOSE_PAIR_DISTANCE = 0.5 + + +def _grid(spacing: float) -> np.ndarray: + """Return a 3x3x3 grid of ``NATOMS`` positions with the given spacing.""" + axis = np.arange(3) * spacing + return np.vstack([g.ravel() for g in np.meshgrid(*([axis] * 3), indexing="ij")]).T + +def _write_system(path: Path, coords: np.ndarray, boxes: np.ndarray) -> None: + """Write frames as a multi-set system in the DeePMD npy layout.""" + nframes = coords.shape[0] path.mkdir(parents=True, exist_ok=True) (path / "type.raw").write_text("\n".join(["0"] * NATOMS) + "\n") (path / "type_map.raw").write_text("TYPE\n") - for index, start in enumerate(range(0, NFRAMES, SET_SIZE)): + for index, start in enumerate(range(0, nframes, SET_SIZE)): subset = path / f"set.{index:03d}" subset.mkdir() frames = slice(start, start + SET_SIZE) - np.save(subset / "coord.npy", coord[frames]) - np.save(subset / "box.npy", box[frames]) - np.save(subset / "energy.npy", np.zeros(SET_SIZE)) - np.save(subset / "force.npy", np.zeros((SET_SIZE, NATOMS * 3))) + nsub = coords[frames].shape[0] + np.save(subset / "coord.npy", coords[frames].reshape(nsub, NATOMS * 3)) + np.save(subset / "box.npy", boxes[frames]) + np.save(subset / "energy.npy", np.zeros(nsub)) + np.save(subset / "force.npy", np.zeros((nsub, NATOMS * 3))) + + +def _box(spacing: float) -> np.ndarray: + return (np.eye(3) * spacing * 3).ravel() + + +def _uniform_frames() -> tuple[np.ndarray, np.ndarray]: + """Identical frames, so the statistics are known analytically.""" + coords = np.repeat(_grid(UNIFORM_SPACING)[np.newaxis], NFRAMES, axis=0) + boxes = np.tile(_box(UNIFORM_SPACING), (NFRAMES, 1)) + return coords, boxes + + +def _mixed_frames() -> tuple[np.ndarray, np.ndarray]: + """Frames whose two extrema are contributed by different sets. + + The first set holds a frame with one pair pulled close, which owns the + global minimum distance but barely changes any neighbor count. The last + set holds a denser frame, whose second shell falls inside the cutoff and + owns the global maximum neighbor count. A reduction that keeps only one + batch therefore cannot report both. + """ + sparse = _grid(SPARSE_SPACING) + close_pair = sparse.copy() + # index 9 is (spacing, 0, 0); move atom 0 to sit CLOSE_PAIR_DISTANCE from it + close_pair[0] = sparse[9] - np.array([CLOSE_PAIR_DISTANCE, 0.0, 0.0]) + dense = _grid(DENSE_SPACING) + + coords = np.repeat(sparse[np.newaxis], NFRAMES, axis=0) + boxes = np.tile(_box(SPARSE_SPACING), (NFRAMES, 1)) + coords[0] = close_pair + coords[-1] = dense + boxes[-1] = _box(DENSE_SPACING) + return coords, boxes def _stat_with_batch_size( - monkeypatch: Any, system: str, batch_atoms: int + monkeypatch: Any, system: str, rcut: float, batch_atoms: int ) -> tuple[tuple[float, np.ndarray], list[int]]: """Run the statistics at a pinned batch size. @@ -87,52 +131,75 @@ def spy( systems=[system], batch_size=1, test_size=1, - rcut=RCUT, + rcut=rcut, type_map=["TYPE"], ) data.get_batch() - stat = NeighborStat(data.get_ntypes(), RCUT, mixed_type=False).get_stat(data) + stat = NeighborStat(data.get_ntypes(), rcut, mixed_type=False).get_stat(data) return stat, frames_per_call +def _single_frame_stat( + monkeypatch: Any, path: Path, coord: np.ndarray, box: np.ndarray +) -> tuple[float, np.ndarray]: + """Statistics of one frame on its own, as a reference extremum.""" + _write_system(path, coord[np.newaxis], box[np.newaxis]) + stat, _ = _stat_with_batch_size(monkeypatch, str(path), MIXED_RCUT, NATOMS) + return stat + + def test_batching_splits_the_work_and_keeps_the_result( monkeypatch: Any, tmp_path: Path ) -> None: """Batching must change how the work is split, not what it produces. Sending a whole set at once allocates an ``[nframes, nloc, nall, 3]`` - tensor, which is why the statistics are batched. Comparing the results - alone would pass without any batching, so the split itself is asserted. + tensor, which is why the statistics are batched. Both the split and the + reduction across batches are asserted: the frames are built so that the + two extrema come from different sets, so keeping only one batch would + get one of them wrong. """ + coords, boxes = _mixed_frames() system = tmp_path / "system" - _write_system(system) + _write_system(system, coords, boxes) whole, whole_calls = _stat_with_batch_size( - monkeypatch, str(system), NATOMS * NFRAMES + monkeypatch, str(system), MIXED_RCUT, NATOMS * NFRAMES + ) + batched, batched_calls = _stat_with_batch_size( + monkeypatch, str(system), MIXED_RCUT, NATOMS ) - batched, batched_calls = _stat_with_batch_size(monkeypatch, str(system), NATOMS) # one call per set when everything fits, one call per frame when it does not assert whole_calls == [SET_SIZE] * (NFRAMES // SET_SIZE) assert batched_calls == [1] * NFRAMES assert len(batched_calls) > len(whole_calls) - assert batched[0] == whole[0] - np.testing.assert_array_equal(batched[1], whole[1]) + # the frame in the first set owns the shortest distance, the one in the + # last set the largest neighbor count + first = _single_frame_stat(monkeypatch, tmp_path / "first", coords[0], boxes[0]) + last = _single_frame_stat(monkeypatch, tmp_path / "last", coords[-1], boxes[-1]) + assert first[0] < last[0] + assert first[1].max() < last[1].max() + + for stat in (whole, batched): + assert stat[0] == first[0] + np.testing.assert_array_equal(stat[1], last[1]) def test_batched_statistics_are_correct(monkeypatch: Any, tmp_path: Path) -> None: """The reduction over batches reproduces the known geometry. - On a unit grid the shortest distance is 1.0, and within ``rcut`` each - atom sees its six axis-aligned images. + On a unit grid the shortest distance is 1.0, and just past it each atom + sees its six axis-aligned images. """ + coords, boxes = _uniform_frames() system = tmp_path / "system" - _write_system(system) + _write_system(system, coords, boxes) (min_nbor_dist, max_nbor_size), _ = _stat_with_batch_size( - monkeypatch, str(system), NATOMS + monkeypatch, str(system), UNIFORM_RCUT, NATOMS ) - assert min_nbor_dist == 1.0 + assert min_nbor_dist == UNIFORM_SPACING np.testing.assert_array_equal(max_nbor_size, [6]) From 22cedf5fec66e0c98d4c4e66d32e56139a2a4ca8 Mon Sep 17 00:00:00 2001 From: yckbz Date: Sun, 18 Oct 2026 23:01:49 +0800 Subject: [PATCH 15/15] test(consistent): gate the pt_expt leg on INSTALLED_PT_EXPT INSTALLED_PT only says torch imports; the backend itself can still be absent, and the entry point then raises instead of the test skipping. --- source/tests/consistent/test_neighbor_stat.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/source/tests/consistent/test_neighbor_stat.py b/source/tests/consistent/test_neighbor_stat.py index 35ff97ae37..34c6ca0e22 100644 --- a/source/tests/consistent/test_neighbor_stat.py +++ b/source/tests/consistent/test_neighbor_stat.py @@ -16,6 +16,7 @@ INSTALLED_JAX, INSTALLED_PD, INSTALLED_PT, + INSTALLED_PT_EXPT, INSTALLED_TF, ) @@ -82,7 +83,7 @@ def test_neighbor_stat_tf(self) -> None: def test_neighbor_stat_pt(self) -> None: self.run_neighbor_stat("pytorch") - @unittest.skipUnless(INSTALLED_PT, "pytorch is not installed") + @unittest.skipUnless(INSTALLED_PT_EXPT, "PyTorch Exportable is not installed") def test_neighbor_stat_pt_expt(self) -> None: self.run_neighbor_stat("pt-expt")