From 1171c9f038a6ba001f8cdb54901d150c8476ac8f Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 8 Sep 2026 16:55:00 +0800
Subject: [PATCH 1/5] feat(pt,pd): exact full-scan output statistics via
 streaming QR

Output statistics were fitted from `data_stat_nbatch` batches per system.
Elements too rare to appear in that sample got no support in the least-squares
design matrix, so their bias collapsed to the min-norm zero and they were also
missing from `observed_type`.

Add `model.data_stat_full` (default false), which computes bias, std and
observed types from every training frame. `ReduStatAccumulator` folds frames
into a running QR factor of `[natoms | output_redu | 1]`; since `R.T @ R ==
A.T @ A`, the solution matches a fit over all frames while memory stays at
O((ntypes + ndim + 1)^2). `assigned_bias`, `intensive` and `atom_exclude_types`
are column-linear, so they apply exactly to the compressed factor.

The trainer attaches a `ReduStatScanner` to the stat sampler; the wrapped
sampler forwards it and `_compute_output_stats_global` picks it up. The scan
rebuilds each loader without its sampler so a distributed or weighted sampler
cannot hide frames from the chief. Delta bias and `stats_distinguish_types=False`
cannot be expressed on the compressed factor and fall back to the sampled path.
Descriptor input statistics still use `data_stat_nbatch`.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ByhbNb5V8vZAGqY88H3wF6
---
 .../dpmodel/atomic_model/base_atomic_model.py |  15 +-
 deepmd/dpmodel/utils/stat.py                  |  26 +++
 .../pd/model/atomic_model/dp_atomic_model.py  |   6 +
 deepmd/pd/train/training.py                   |  18 ++
 deepmd/pd/utils/stat.py                       | 143 +++++++++++-
 .../model/atomic_model/base_atomic_model.py   |  15 +-
 deepmd/pt/train/training.py                   |  22 ++
 deepmd/pt/utils/stat.py                       | 153 ++++++++++++-
 deepmd/utils/argcheck.py                      |  14 ++
 deepmd/utils/out_stat.py                      | 215 ++++++++++++++++++
 source/tests/common/test_out_stat.py          | 114 ++++++++++
 11 files changed, 728 insertions(+), 13 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py
index 7050a4a62e..6587fe41ae 100644
--- a/deepmd/dpmodel/atomic_model/base_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py
@@ -106,6 +106,7 @@ def _collect_and_set_observed_type(
             _restore_observed_type_from_file,
             _save_observed_type_to_file,
             collect_observed_types,
+            observed_types_from_counts,
         )
 
         if preset_observed_type is not None:
@@ -113,8 +114,13 @@ def _collect_and_set_observed_type(
         else:
             observed = _restore_observed_type_from_file(stat_file_path)
             if observed is None:
-                sampled = sampled_func()
-                observed = collect_observed_types(sampled, self.type_map)
+                scanner = getattr(sampled_func, "redu_stat_scanner", None)
+                if scanner is not None:
+                    observed = observed_types_from_counts(
+                        scanner.natoms_total(len(self.type_map)), self.type_map
+                    )
+                else:
+                    observed = collect_observed_types(sampled_func(), self.type_map)
                 _save_observed_type_to_file(stat_file_path, observed)
             self._observed_type = observed
 
@@ -767,6 +773,11 @@ def wrapped_sampler() -> list[dict]:
                             sample["find_fparam"] = np.bool_(True)
             return sampled
 
+        # the full-data scanner, when the trainer attached one, is part of the
+        # sampler contract and must survive wrapping
+        wrapped_sampler.redu_stat_scanner = getattr(
+            sampled_func, "redu_stat_scanner", None
+        )
         return wrapped_sampler
 
     def change_out_bias(
diff --git a/deepmd/dpmodel/utils/stat.py b/deepmd/dpmodel/utils/stat.py
index 900e101b7d..f62e80defe 100644
--- a/deepmd/dpmodel/utils/stat.py
+++ b/deepmd/dpmodel/utils/stat.py
@@ -65,6 +65,32 @@ def collect_observed_types(sampled: list[dict], type_map: list[str]) -> list[str
     return sort_element_type(observed_types)
 
 
+def observed_types_from_counts(
+    natoms_total: np.ndarray, type_map: list[str]
+) -> list[str]:
+    """Collect observed element types from per-type atom counts.
+
+    Parameters
+    ----------
+    natoms_total : np.ndarray
+        Total occurrences of each type, shape ``[ntypes]``.
+    type_map : list[str]
+        Mapping from type index to element symbol.
+
+    Returns
+    -------
+    list[str]
+        Sorted list of observed element symbols.
+    """
+    from deepmd.utils.econf_embd import (
+        sort_element_type,
+    )
+
+    return sort_element_type(
+        [type_map[i] for i in np.flatnonzero(np.asarray(natoms_total) > 0)]
+    )
+
+
 def _restore_observed_type_from_file(
     stat_file_path: DPPath | None,
 ) -> list[str] | None:
diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py
index 334ebdc5f5..2dca9ea65f 100644
--- a/deepmd/pd/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py
@@ -409,6 +409,12 @@ def wrapped_sampler() -> list[dict]:
                     sample["atom_exclude_types"] = list(atom_exclude_types)
             return sampled
 
+        # the full-data scanner, when the trainer attached one, is part of the
+        # sampler contract and must survive wrapping
+        wrapped_sampler.redu_stat_scanner = getattr(
+            sampled_func, "redu_stat_scanner", None
+        )
+
         self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path)
         self.compute_fitting_input_stat(wrapped_sampler, stat_file_path)
         if compute_or_load_out_stat:
diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py
index 688f0c81c0..a33908c632 100644
--- a/deepmd/pd/train/training.py
+++ b/deepmd/pd/train/training.py
@@ -73,6 +73,7 @@
 )
 from deepmd.pd.utils.stat import (
     make_stat_input,
+    scan_redu_stats,
 )
 from deepmd.pd.utils.utils import (
     nvprof_context,
@@ -84,6 +85,9 @@
 from deepmd.utils.finetune import (
     warn_configuration_mismatch_during_finetune,
 )
+from deepmd.utils.out_stat import (
+    ReduStatScanner,
+)
 from deepmd.utils.path import (
     DPH5Path,
 )
@@ -233,6 +237,7 @@ def single_model_stat(
             _validation_data: Any | None,
             _stat_file_path: str | Path | None,
             _data_requirement: list[DataRequirementItem],
+            _data_stat_full: bool = False,
             finetune_has_new_type: bool = False,
         ) -> Any:
             _data_requirement += get_additional_data_requirement(_model)
@@ -249,6 +254,15 @@ def get_sample() -> dict[str, Any]:
                 )
                 return sampled
 
+            if _data_stat_full:
+                # sampling a few batches per system can miss rare elements
+                # entirely; scan every frame for the output statistics instead
+                get_sample.redu_stat_scanner = ReduStatScanner(
+                    lambda ntypes, keys, intensive: scan_redu_stats(
+                        _training_data.dataloaders, ntypes, keys, intensive=intensive
+                    )
+                )
+
             if (not resuming or finetune_has_new_type) and self.rank == 0:
                 _model.compute_or_load_stat(
                     sampled_func=get_sample,
@@ -310,6 +324,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                 validation_data,
                 stat_file_path,
                 self.loss.label_requirement,
+                _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
                 else False,
@@ -349,6 +364,9 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                     validation_data[model_key],
                     stat_file_path[model_key],
                     self.loss[model_key].label_requirement,
+                    _data_stat_full=model_params["model_dict"][model_key].get(
+                        "data_stat_full", False
+                    ),
                     finetune_has_new_type=self.finetune_links[
                         model_key
                     ].get_has_new_type()
diff --git a/deepmd/pd/utils/stat.py b/deepmd/pd/utils/stat.py
index b80af7b60e..1af3fd0933 100644
--- a/deepmd/pd/utils/stat.py
+++ b/deepmd/pd/utils/stat.py
@@ -5,6 +5,7 @@
 )
 from collections.abc import (
     Callable,
+    Sequence,
 )
 from typing import (
     Any,
@@ -12,6 +13,10 @@
 
 import numpy as np
 import paddle
+from paddle.io import (
+    BatchSampler,
+    DataLoader,
+)
 
 from deepmd.pd.utils import (
     AtomExcludeMask,
@@ -25,6 +30,9 @@
     to_paddle_tensor,
 )
 from deepmd.utils.out_stat import (
+    ReduScanResult,
+    ReduStatAccumulator,
+    ReduStatScanner,
     compute_stats_do_not_distinguish_types,
     compute_stats_from_atomic,
     compute_stats_from_redu,
@@ -94,6 +102,88 @@ def make_stat_input(
     return lst
 
 
+def _full_pass_loader(dataloader: Any) -> Any:
+    """Return a loader that covers the dataset once, whatever sampler it uses.
+
+    Training loaders may carry a distributed or weighted sampler, which would
+    hide part of the data from a scan that is meant to be exhaustive.
+    """
+    sampler = getattr(dataloader, "batch_sampler", None)
+    if sampler is None or getattr(sampler, "batch_size", None) is None:
+        return dataloader
+    return DataLoader(
+        dataset=dataloader.dataset,
+        num_workers=0,
+        batch_sampler=BatchSampler(
+            dataloader.dataset, shuffle=False, batch_size=sampler.batch_size
+        ),
+        collate_fn=dataloader.collate_fn,
+        use_buffer_reader=False,
+        places=["cpu"],
+    )
+
+
+def scan_redu_stats(
+    dataloaders: list[Any],
+    ntypes: int,
+    keys: Sequence[str],
+    intensive: bool = False,
+) -> ReduScanResult:
+    """Accumulate exact reduced-label statistics over every training frame.
+
+    Where :func:`make_stat_input` keeps a few batches per system, this scans the
+    whole training set but retains only per-type atom counts and reduced labels,
+    compressed into one :class:`ReduStatAccumulator` per key. Elements that are
+    too rare to survive batch sampling therefore still enter the regression, at
+    a memory cost that does not grow with the number of frames.
+
+    Parameters
+    ----------
+    dataloaders
+        One data loader for each system.
+    ntypes
+        The number of atom types.
+    keys
+        Output labels whose statistics are accumulated.
+    intensive
+        Whether the fitting target is intensive.
+
+    Returns
+    -------
+    ReduScanResult
+        The accumulators, the per-type atom counts and the frame count.
+    """
+    stats: dict[str, ReduStatAccumulator] = {}
+    natoms_total = np.zeros(ntypes, dtype=np.int64)
+    nframes = 0
+    log.info(f"Scanning all frames of {len(dataloaders)} systems for output statistics")
+    for dataloader in dataloaders:
+        for batch in _full_pass_loader(dataloader):
+            natoms_key = "real_natoms_vec" if "real_natoms_vec" in batch else "natoms"
+            # natoms is [nframes, 2 + ntypes]; the first two are nall/nloc
+            natoms = to_numpy_array(batch[natoms_key])[:, 2:]
+            natoms_total += natoms.sum(axis=0).astype(np.int64)
+            nframes += natoms.shape[0]
+            for key in keys:
+                if key not in batch or float(batch.get(f"find_{key}", 0.0)) <= 0.0:
+                    continue
+                label = to_numpy_array(batch[key])
+                if key not in stats:
+                    var_shape = list(label.shape[1:])
+                    stats[key] = ReduStatAccumulator(
+                        ntypes,
+                        int(np.prod(var_shape)) if var_shape else 1,
+                        var_shape,
+                        intensive=intensive,
+                    )
+                stats[key].add(label, natoms)
+    log.info(
+        f"Scanned {nframes} frames; "
+        f"{int(np.count_nonzero(natoms_total))} of {ntypes} types observed"
+    )
+    return ReduScanResult(stats=stats, natoms_total=natoms_total, nframes=nframes)
+
+
 def _restore_from_file(
     stat_file_path: DPPath,
     keys: list[str] = ["energy"],
@@ -380,6 +470,7 @@ def compute_output_stats(
             stats_distinguish_types,
             intensive,
             model_pred_g,
+            getattr(merged, "redu_stat_scanner", None),
         )
         bias_atom_a, std_atom_a = _compute_output_stats_atomic(
             sampled,
@@ -426,8 +517,14 @@ def _compute_output_stats_global(
     stats_distinguish_types: bool = True,
     intensive: bool = False,
     model_pred: dict[str, np.ndarray] | None = None,
+    redu_scanner: ReduStatScanner | None = None,
 ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
-    """This function only handle stat computation from reduced global labels."""
+    """This function only handle stat computation from reduced global labels.
+
+    When *redu_scanner* is given, the bias and std come from a scan of every
+    training frame instead of the sampled batches, which keeps rare elements
+    from being under-represented in the regression.
+    """
     # return directly if no global samples
     if global_sampled_idx is None or all(
         len(v) == 0 for v in global_sampled_idx.values()
@@ -487,10 +584,35 @@ def _compute_output_stats_global(
             if kk in merged_output
         }
 
+    scan = None
+    if redu_scanner is not None:
+        if model_pred is not None or not stats_distinguish_types:
+            log.warning(
+                "Falling back to sampled output statistics: a full scan supports "
+                "neither delta bias nor stats_distinguish_types=False."
+            )
+        else:
+            scan = redu_scanner.scan(ntypes, keys, intensive)
+    type_mask = (
+        to_numpy_array(
+            AtomExcludeMask(ntypes, sampled[0]["atom_exclude_types"]).get_type_mask()
+        )
+        if scan is not None and "atom_exclude_types" in sampled[0]
+        else None
+    )
+
     bias_atom_e = {}
     std_atom_e = {}
+    scanned_keys = set()
     for kk in keys:
-        if kk in stats_input:
+        if scan is not None and kk in scan.stats:
+            bias_atom_e[kk], std_atom_e[kk] = scan.stats[kk].solve(
+                assigned_bias=assigned_atom_ener[kk],
+                rcond=rcond,
+                type_mask=type_mask,
+            )
+            scanned_keys.add(kk)
+        elif kk in stats_input:
             if not stats_distinguish_types:
                 bias_atom_e[kk], std_atom_e[kk] = (
                     compute_stats_do_not_distinguish_types(
@@ -514,23 +636,32 @@ def _compute_output_stats_global(
 
     # unbias_e is only used for print rmse
 
+    sampled_keys = [
+        kk for kk in bias_atom_e if kk not in scanned_keys and kk in merged_natoms
+    ]
     if model_pred is None:
         unbias_e = {
             kk: merged_natoms[kk] @ bias_atom_e[kk].reshape([ntypes, -1])
-            for kk in bias_atom_e.keys()
+            for kk in sampled_keys
         }
     else:
         unbias_e = {
             kk: model_pred[kk].reshape([nf[kk], -1])
             + merged_natoms[kk] @ bias_atom_e[kk].reshape([ntypes, -1])
-            for kk in bias_atom_e.keys()
+            for kk in sampled_keys
         }
-    atom_numbs = {kk: merged_natoms[kk].sum(-1) for kk in bias_atom_e.keys()}
+    atom_numbs = {kk: merged_natoms[kk].sum(-1) for kk in sampled_keys}
 
     def rmse(x: np.ndarray) -> float:
         return np.sqrt(np.mean(np.square(x)))
 
-    for kk in bias_atom_e.keys():
+    for kk in scanned_keys:
+        log.info(
+            f"Std of {kk} residual after linear regression over "
+            f"{scan.stats[kk].nframes} frames is: {std_atom_e[kk].reshape(-1)[0]} "
+            f"in the unit of {kk}."
+        )
+    for kk in sampled_keys:
         rmse_ae = rmse(
             (
                 unbias_e[kk].reshape([nf[kk], -1])
diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py
index d9689055da..3da0e16425 100644
--- a/deepmd/pt/model/atomic_model/base_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/base_atomic_model.py
@@ -119,6 +119,7 @@ def _collect_and_set_observed_type(
             _restore_observed_type_from_file,
             _save_observed_type_to_file,
             collect_observed_types,
+            observed_types_from_counts,
         )
 
         if preset_observed_type is not None:
@@ -126,8 +127,13 @@ def _collect_and_set_observed_type(
         else:
             observed = _restore_observed_type_from_file(stat_file_path)
             if observed is None:
-                sampled = sampled_func()
-                observed = collect_observed_types(sampled, self.type_map)
+                scanner = getattr(sampled_func, "redu_stat_scanner", None)
+                if scanner is not None:
+                    observed = observed_types_from_counts(
+                        scanner.natoms_total(len(self.type_map)), self.type_map
+                    )
+                else:
+                    observed = collect_observed_types(sampled_func(), self.type_map)
                 _save_observed_type_to_file(stat_file_path, observed)
             self._observed_type = observed
 
@@ -254,6 +260,11 @@ def wrapped_sampler() -> list[dict]:
                         sample["fparam"] = default_fparam.repeat(nframe, 1)
             return sampled
 
+        # the full-data scanner, when the trainer attached one, is part of the
+        # sampler contract and must survive wrapping
+        wrapped_sampler.redu_stat_scanner = getattr(
+            sampled_func, "redu_stat_scanner", None
+        )
         return wrapped_sampler
 
     def reinit_atom_exclude(
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index dd515b1a64..40e1659f95 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -98,6 +98,7 @@
 from deepmd.pt.utils.stat import (
     make_stat_input,
     min_pair_dist_frame_mask,
+    scan_redu_stats,
     select_batch_frames,
 )
 from deepmd.pt.utils.utils import (
@@ -129,6 +130,9 @@
 from deepmd.utils.finetune import (
     warn_configuration_mismatch_during_finetune,
 )
+from deepmd.utils.out_stat import (
+    ReduStatScanner,
+)
 
 if torch.__version__.startswith("2"):
     import torch._dynamo
@@ -395,6 +399,7 @@ def single_model_stat(
             _training_data: DpLoaderSet,
             _stat_file_spec: StatFileSpec,
             _min_pair_dist: float = 0.0,
+            _data_stat_full: bool = False,
             finetune_has_new_type: bool = False,
             preset_observed_type: list[str] | None = None,
         ) -> Callable[[], Any]:
@@ -408,6 +413,19 @@ def get_sample() -> Any:
                 )
                 return sampled
 
+            if _data_stat_full:
+                # sampling a few batches per system can miss rare elements
+                # entirely; scan every frame for the output statistics instead
+                get_sample.redu_stat_scanner = ReduStatScanner(
+                    lambda ntypes, keys, intensive: scan_redu_stats(
+                        _training_data.dataloaders,
+                        ntypes,
+                        keys,
+                        intensive=intensive,
+                        min_pair_dist=_min_pair_dist,
+                    )
+                )
+
             if not has_initial_state or finetune_has_new_type:
 
                 def initialize_statistics() -> None:
@@ -529,6 +547,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                 training_data,
                 self.stat_file_specs["Default"],
                 _min_pair_dist=min_pair_dist,
+                _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
                 else False,
@@ -613,6 +632,9 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                     training_data[model_key],
                     self.stat_file_specs[model_key],
                     _min_pair_dist=min_pair_dist,
+                    _data_stat_full=model_params["model_dict"][model_key].get(
+                        "data_stat_full", False
+                    ),
                     finetune_has_new_type=self.finetune_links[
                         model_key
                     ].get_has_new_type()
diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
index 5b245d187b..fdf7f87f09 100644
--- a/deepmd/pt/utils/stat.py
+++ b/deepmd/pt/utils/stat.py
@@ -5,6 +5,7 @@
 )
 from collections.abc import (
     Callable,
+    Sequence,
 )
 from typing import (
     Any,
@@ -12,6 +13,9 @@
 
 import numpy as np
 import torch
+from torch.utils.data import (
+    DataLoader,
+)
 
 from deepmd.pt.utils import (
     AtomExcludeMask,
@@ -25,6 +29,9 @@
     to_torch_tensor,
 )
 from deepmd.utils.out_stat import (
+    ReduScanResult,
+    ReduStatAccumulator,
+    ReduStatScanner,
     compute_stats_do_not_distinguish_types,
     compute_stats_from_atomic,
     compute_stats_from_redu,
@@ -44,6 +51,7 @@
     _restore_observed_type_from_file,
     _save_observed_type_to_file,
     collect_observed_types,
+    observed_types_from_counts,
 )
 
 __all__ = [
@@ -51,6 +59,8 @@
     "_save_observed_type_to_file",
     "collect_observed_types",
     "min_pair_dist_frame_mask",
+    "observed_types_from_counts",
+    "scan_redu_stats",
     "select_batch_frames",
 ]
 
@@ -222,6 +232,103 @@ def make_stat_input(
     return lst
 
 
+def _full_pass_loader(dataloader: Any) -> Any:
+    """Return a loader that covers the dataset once, whatever sampler it uses.
+
+    Training loaders may carry a distributed or weighted sampler, which would
+    hide part of the data from a scan that is meant to be exhaustive.
+    """
+    if dataloader.batch_size is None:
+        # a custom batch sampler owns the batching; leave it alone
+        return dataloader
+    with torch.device("cpu"):
+        return DataLoader(
+            dataloader.dataset,
+            batch_size=dataloader.batch_size,
+            shuffle=False,
+            num_workers=0,
+            drop_last=False,
+            collate_fn=dataloader.collate_fn,
+        )
+
+
+def scan_redu_stats(
+    dataloaders: list[Any],
+    ntypes: int,
+    keys: Sequence[str],
+    intensive: bool = False,
+    min_pair_dist: float = 0.0,
+) -> ReduScanResult:
+    """Accumulate exact reduced-label statistics over every training frame.
+
+    Where :func:`make_stat_input` keeps a few batches per system, this scans the
+    whole training set but retains only per-type atom counts and reduced labels,
+    compressed into one :class:`ReduStatAccumulator` per key. Elements that are
+    too rare to survive batch sampling therefore still enter the regression, at
+    a memory cost that does not grow with the number of frames.
+
+    Parameters
+    ----------
+    dataloaders
+        One data loader for each system.
+    ntypes
+        The number of atom types.
+    keys
+        Output labels whose statistics are accumulated.
+    intensive
+        Whether the fitting target is intensive.
+    min_pair_dist
+        Minimum allowed pair distance in Angstrom. Frames below the threshold
+        are excluded.
+
+    Returns
+    -------
+    ReduScanResult
+        The accumulators, the per-type atom counts and the frame count.
+    """
+    stats: dict[str, ReduStatAccumulator] = {}
+    natoms_total = np.zeros(ntypes, dtype=np.int64)
+    nframes = 0
+    log.info(
+        "Scanning all frames of %d systems for output statistics", len(dataloaders)
+    )
+    with torch.device("cpu"):
+        for dataloader in dataloaders:
+            for batch in _full_pass_loader(dataloader):
+                frame_mask = min_pair_dist_frame_mask(batch, min_pair_dist)
+                if frame_mask is not None:
+                    if not torch.any(frame_mask):
+                        continue
+                    batch = select_batch_frames(batch, frame_mask)
+                natoms_key = (
+                    "real_natoms_vec" if "real_natoms_vec" in batch else "natoms"
+                )
+                # natoms is [nframes, 2 + ntypes]; the first two are nall/nloc
+                natoms = to_numpy_array(batch[natoms_key])[:, 2:]
+                natoms_total += natoms.sum(axis=0).astype(np.int64)
+                nframes += natoms.shape[0]
+                for key in keys:
+                    if key not in batch or float(batch.get(f"find_{key}", 0.0)) <= 0.0:
+                        continue
+                    label = to_numpy_array(batch[key])
+                    if key not in stats:
+                        var_shape = list(label.shape[1:])
+                        stats[key] = ReduStatAccumulator(
+                            ntypes,
+                            int(np.prod(var_shape)) if var_shape else 1,
+                            var_shape,
+                            intensive=intensive,
+                        )
+                    stats[key].add(label, natoms)
+    log.info(
+        "Scanned %d frames; %d of %d types observed",
+        nframes,
+        int(np.count_nonzero(natoms_total)),
+        ntypes,
+    )
+    return ReduScanResult(stats=stats, natoms_total=natoms_total, nframes=nframes)
+
+
 def _restore_from_file(
     stat_file_path: DPPath | None,
     keys: list[str],
@@ -565,6 +672,7 @@ def compute_output_stats(
             stats_distinguish_types,
             intensive,
             model_pred_g,
+            getattr(merged, "redu_stat_scanner", None),
         )
         bias_atom_a, std_atom_a = _compute_output_stats_atomic(
             sampled,
@@ -616,8 +724,14 @@ def _compute_output_stats_global(
     stats_distinguish_types: bool = True,
     intensive: bool = False,
     model_pred: dict[str, np.ndarray] | None = None,
+    redu_scanner: ReduStatScanner | None = None,
 ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
-    """This function only handle stat computation from reduced global labels."""
+    """This function only handle stat computation from reduced global labels.
+
+    When *redu_scanner* is given, the bias and std come from a scan of every
+    training frame instead of the sampled batches, which keeps rare elements
+    from being under-represented in the regression.
+    """
     # return directly if no global samples
     if global_sampled_idx is None or all(
         len(v) == 0 for v in global_sampled_idx.values()
@@ -677,10 +791,35 @@ def _compute_output_stats_global(
             if kk in merged_output
         }
 
+    scan = None
+    if redu_scanner is not None:
+        if model_pred is not None or not stats_distinguish_types:
+            log.warning(
+                "Falling back to sampled output statistics: a full scan supports "
+                "neither delta bias nor stats_distinguish_types=False."
+            )
+        else:
+            scan = redu_scanner.scan(ntypes, keys, intensive)
+    type_mask = (
+        to_numpy_array(
+            AtomExcludeMask(ntypes, sampled[0]["atom_exclude_types"]).get_type_mask()
+        )
+        if scan is not None and "atom_exclude_types" in sampled[0]
+        else None
+    )
+
     bias_atom_e = {}
     std_atom_e = {}
+    scanned_keys = set()
     for kk in keys:
-        if kk in stats_input:
+        if scan is not None and kk in scan.stats:
+            bias_atom_e[kk], std_atom_e[kk] = scan.stats[kk].solve(
+                assigned_bias=assigned_atom_ener[kk],
+                rcond=rcond,
+                type_mask=type_mask,
+            )
+            scanned_keys.add(kk)
+        elif kk in stats_input:
             if not stats_distinguish_types:
                 bias_atom_e[kk], std_atom_e[kk] = (
                     compute_stats_do_not_distinguish_types(
@@ -707,6 +846,8 @@ def _compute_output_stats_global(
 
     unbias_e = {}
     for kk in bias_atom_e.keys():
+        if kk in scanned_keys or kk not in merged_natoms:
+            continue
         coeffs = merged_natoms[kk]
         if intensive:
             total_atoms = coeffs.sum(axis=1, keepdims=True)
@@ -719,7 +860,13 @@ def _compute_output_stats_global(
     def rmse(x: np.ndarray) -> float:
         return np.sqrt(np.mean(np.square(x)))
 
-    for kk in bias_atom_e.keys():
+    for kk in scanned_keys:
+        log.info(
+            f"Std of {kk} residual after linear regression over "
+            f"{scan.stats[kk].nframes} frames is: {std_atom_e[kk].reshape(-1)[0]} "
+            f"in the unit of {kk}."
+        )
+    for kk in unbias_e.keys():
         diff = unbias_e[kk].reshape(nf[kk], -1) - merged_output[kk].reshape(nf[kk], -1)
         if not intensive:
             diff /= merged_natoms[kk].sum(axis=-1, keepdims=True)
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 65ab8a85a2..deca3b8374 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -3398,6 +3398,13 @@ def model_args(
     doc_type_map = "A list of strings. Give the name to each type of atoms. It is noted that the number of atom type of training system must be less than 128 in a GPU environment. If not given, type.raw in each system should use the same type indexes, and type_map.raw will take no effect."
     doc_data_stat_nbatch = "The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics."
     doc_data_stat_protect = "Protect parameter for atomic energy regression."
+    doc_data_stat_full = (
+        "Scan every frame of the training data to compute the output statistics "
+        "(bias and standard deviation of the fitting target) exactly, instead of "
+        "estimating them from `data_stat_nbatch` batches per system. Recommended "
+        "for datasets containing rare elements, whose bias is otherwise fitted "
+        "from too few frames. Input statistics still use `data_stat_nbatch`."
+    )
     doc_data_bias_nsample = "The number of training samples in a system to compute and change the energy bias."
     doc_type_embedding = "The type embedding. In other backends, the type embedding is already included in the descriptor."
     doc_modifier = "The modifier of model output."
@@ -3438,6 +3445,13 @@ def model_args(
                 default=1e-2,
                 doc=doc_data_stat_protect,
             ),
+            Argument(
+                "data_stat_full",
+                bool,
+                optional=True,
+                default=False,
+                doc=supported_backends("pt", "pd") + doc_data_stat_full,
+            ),
             Argument(
                 "data_bias_nsample",
                 int,
diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py
index ef9db35510..d62d681456 100644
--- a/deepmd/utils/out_stat.py
+++ b/deepmd/utils/out_stat.py
@@ -1,6 +1,14 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Output statistics."""
 
+from collections.abc import (
+    Callable,
+    Sequence,
+)
+from dataclasses import (
+    dataclass,
+)
+
 import numpy as np
 
 from deepmd.env import (
@@ -192,3 +200,210 @@ def compute_stats_do_not_distinguish_types(
     output_std = np.tile(output_std, (computed_output_bias.shape[0], 1))
 
     return computed_output_bias, output_std
+
+
+class ReduStatAccumulator:
+    """Streaming, exact form of :func:`compute_stats_from_redu`.
+
+    Frames are folded into a running QR factor of the augmented design matrix
+    ``[natoms | output_redu | 1]``. Because ``R.T @ R == A.T @ A``, the least
+    squares solution and the residual std are identical to the ones a single
+    call on all frames would return, while memory stays at
+    ``O((ntypes + ndim + 1) ** 2)`` regardless of the number of frames.
+
+    Parameters
+    ----------
+    ntypes
+        The number of atom types.
+    ndim
+        The flattened output dimension.
+    var_shape
+        The unflattened output shape, ``(ndim,)`` when not given.
+    intensive
+        Whether the output is intensive or extensive.
+    """
+
+    def __init__(
+        self,
+        ntypes: int,
+        ndim: int,
+        var_shape: list[int] | None = None,
+        intensive: bool = False,
+    ) -> None:
+        self.ntypes = ntypes
+        self.ndim = ndim
+        self.var_shape = [ndim] if var_shape is None else list(var_shape)
+        self.intensive = intensive
+        self.nframes = 0
+        # total occurrences of each type over every accumulated frame
+        self.natoms_total = np.zeros(ntypes, dtype=np.int64)
+        self._ncols = ntypes + ndim + 1
+        self._r_factor = np.zeros((0, self._ncols), dtype=np.float64)
+        # buffer whole blocks so that one QR amortizes over many small batches
+        self._pending: list[np.ndarray] = []
+        self._pending_rows = 0
+        self._compress_every = max(1024, 8 * self._ncols)
+
+    def add(self, output_redu: np.ndarray, natoms: np.ndarray) -> None:
+        """Accumulate one chunk of frames.
+
+        Parameters
+        ----------
+        output_redu
+            The reduced output value, shape is [nframes, *(odim0, odim1, ...)].
+        natoms
+            The number of atoms of each type, shape is [nframes, ntypes].
+        """
+        natoms = np.asarray(natoms, dtype=np.float64).reshape(-1, self.ntypes)
+        nf = natoms.shape[0]
+        if nf == 0:
+            return
+        output_redu = np.asarray(output_redu, dtype=np.float64).reshape(nf, self.ndim)
+        self.natoms_total += np.rint(natoms.sum(axis=0)).astype(np.int64)
+        self.nframes += nf
+        if self.intensive:
+            natoms = natoms / np.sum(natoms, axis=1, keepdims=True)
+        self._pending.append(
+            np.concatenate(
+                [natoms, output_redu, np.ones((nf, 1), dtype=np.float64)], axis=1
+            )
+        )
+        self._pending_rows += nf
+        if self._pending_rows >= self._compress_every:
+            self._compress()
+
+    def _compress(self) -> None:
+        """Fold the buffered blocks into the running QR factor."""
+        if not self._pending:
+            return
+        self._r_factor = np.linalg.qr(
+            np.concatenate([self._r_factor, *self._pending], axis=0), mode="r"
+        )
+        self._pending.clear()
+        self._pending_rows = 0
+
+    def solve(
+        self,
+        assigned_bias: np.ndarray | None = None,
+        rcond: float | None = None,
+        type_mask: np.ndarray | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Solve the accumulated regression.
+
+        Parameters
+        ----------
+        assigned_bias
+            The assigned output bias, shape is [ntypes, *(odim0, odim1, ...)].
+            Set to a tensor of shape (odim0, odim1, ...) filled with nan if the
+            bias of the type is not assigned.
+        rcond
+            Cut-off ratio for small singular values of a. ``None`` reproduces
+            the numpy default of the equivalent uncompressed problem.
+        type_mask
+            Excluded types, shape is [ntypes]. Types with a zero entry do not
+            contribute to the regression.
+
+        Returns
+        -------
+        np.ndarray
+            The computed output bias, shape is [ntypes, *(odim0, odim1, ...)].
+        np.ndarray
+            The computed output std, shape is [*(odim0, odim1, ...)].
+
+        Raises
+        ------
+        ValueError
+            If no frame has been accumulated.
+        """
+        if self.nframes == 0:
+            raise ValueError("No frame has been accumulated.")
+        self._compress()
+        r_factor = self._r_factor.copy()
+        design = r_factor[:, : self.ntypes]
+        redu = r_factor[:, self.ntypes : self.ntypes + self.ndim]
+        constant = r_factor[:, -1]
+
+        if type_mask is not None:
+            design *= np.asarray(type_mask, dtype=np.float64).reshape(1, self.ntypes)
+
+        assigned_mask = None
+        if assigned_bias is not None:
+            assigned_bias = np.asarray(assigned_bias, dtype=np.float64).reshape(
+                self.ntypes, self.ndim
+            )
+            assigned_mask = ~np.isnan(assigned_bias).any(axis=1)
+            # the same column operations compute_stats_from_redu applies to the
+            # frames; they are linear, so applying them to R is equivalent
+            redu -= design[:, assigned_mask] @ assigned_bias[assigned_mask]
+            design[:, assigned_mask] = 0.0
+
+        if rcond is None:
+            # np.linalg.lstsq scales its default cut-off with the number of
+            # rows, which compression changed; restore the uncompressed one
+            rcond = np.finfo(np.float64).eps * max(self.nframes, self.ntypes)
+        bias, _, _, _ = np.linalg.lstsq(design, redu, rcond=rcond)
+        if assigned_mask is not None:
+            bias[assigned_mask] = assigned_bias[assigned_mask]
+
+        residual = redu - design @ bias
+        residual_mean = (constant @ residual) / self.nframes
+        centered = residual - np.outer(constant, residual_mean)
+        variance = np.sum(centered * centered, axis=0) / self.nframes
+        std = np.sqrt(np.maximum(variance, 0.0))
+        return (
+            bias.reshape([self.ntypes] + self.var_shape),  # noqa: RUF005
+            std.reshape(self.var_shape),
+        )
+
+
+@dataclass
+class ReduScanResult:
+    """Exact statistics collected by one full pass over the training data.
+
+    Parameters
+    ----------
+    stats
+        One accumulator per output key that carries a global label.
+    natoms_total
+        Total occurrences of each type over every scanned frame, shape [ntypes].
+    nframes
+        The number of scanned frames.
+    """
+
+    stats: dict[str, ReduStatAccumulator]
+    natoms_total: np.ndarray
+    nframes: int
+
+
+class ReduStatScanner:
+    """Cache full passes over the training data for a single training run.
+
+    The trainer attaches an instance to the stat sampler; the consumers of that
+    sampler pick it up and use it instead of estimating the output statistics
+    from a handful of sampled batches.
+
+    Parameters
+    ----------
+    scan_fn
+        Backend function performing one pass, called as
+        ``scan_fn(ntypes, keys, intensive)``.
+    """
+
+    def __init__(self, scan_fn: Callable[..., ReduScanResult]) -> None:
+        self._scan_fn = scan_fn
+        self._cache: dict[tuple, ReduScanResult] = {}
+
+    def scan(
+        self, ntypes: int, keys: Sequence[str], intensive: bool = False
+    ) -> ReduScanResult:
+        """Return the statistics for *keys*, scanning the data once per request."""
+        cache_key = (ntypes, tuple(keys), bool(intensive))
+        if cache_key not in self._cache:
+            self._cache[cache_key] = self._scan_fn(ntypes, tuple(keys), bool(intensive))
+        return self._cache[cache_key]
+
+    def natoms_total(self, ntypes: int) -> np.ndarray:
+        """Return the per-type atom counts, reusing any scan already performed."""
+        if self._cache:
+            return next(iter(self._cache.values())).natoms_total
+        return self.scan(ntypes, ()).natoms_total
diff --git a/source/tests/common/test_out_stat.py b/source/tests/common/test_out_stat.py
index 4ec212600a..f6d2b3ab03 100644
--- a/source/tests/common/test_out_stat.py
+++ b/source/tests/common/test_out_stat.py
@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import unittest
+import unittest.mock
 
 import numpy as np
 
 from deepmd.utils.out_stat import (
+    ReduStatAccumulator,
+    ReduStatScanner,
     compute_stats_do_not_distinguish_types,
     compute_stats_from_atomic,
     compute_stats_from_redu,
@@ -203,3 +206,114 @@ def test_compute_stats_from_atomic(self) -> None:
             reference_std,
             rtol=1e-7,
         )
+
+
+class TestReduStatAccumulator(unittest.TestCase):
+    """The streaming accumulator must reproduce compute_stats_from_redu exactly."""
+
+    def setUp(self) -> None:
+        rng = np.random.default_rng(20260908)
+        self.ntypes = 6
+        self.ndim = 3
+        nframes = 500
+        self.natoms = rng.integers(1, 8, size=(nframes, self.ntypes))
+        # a rare element present in a single frame, which batch sampling misses
+        self.natoms[:, 3] = 0
+        self.natoms[0, 3] = 1
+        self.mean = rng.random((self.ntypes, self.ndim)) * 1e3
+        self.output_redu = self.natoms @ self.mean + rng.normal(
+            scale=1e-2, size=(nframes, self.ndim)
+        )
+        return super().setUp()
+
+    def _accumulate(self, intensive: bool = False) -> ReduStatAccumulator:
+        acc = ReduStatAccumulator(
+            self.ntypes, self.ndim, [self.ndim], intensive=intensive
+        )
+        for start in range(0, self.natoms.shape[0], 7):
+            acc.add(self.output_redu[start : start + 7], self.natoms[start : start + 7])
+        return acc
+
+    def test_matches_compute_stats_from_redu(self) -> None:
+        for intensive in (False, True):
+            with self.subTest(intensive=intensive):
+                ref_bias, ref_std = compute_stats_from_redu(
+                    self.output_redu, self.natoms, intensive=intensive
+                )
+                bias, std = self._accumulate(intensive).solve()
+                np.testing.assert_allclose(bias, ref_bias, rtol=1e-9)
+                np.testing.assert_allclose(std, ref_std, rtol=1e-9)
+
+    def test_matches_with_assigned_bias(self) -> None:
+        assigned_bias = np.full((self.ntypes, self.ndim), np.nan)
+        assigned_bias[1] = self.mean[1]
+        assigned_bias[4] = self.mean[4]
+        ref_bias, ref_std = compute_stats_from_redu(
+            self.output_redu.copy(),
+            self.natoms.copy(),
+            assigned_bias=assigned_bias,
+        )
+        bias, std = self._accumulate().solve(assigned_bias=assigned_bias)
+        np.testing.assert_allclose(bias, ref_bias, rtol=1e-9)
+        np.testing.assert_allclose(std, ref_std, rtol=1e-9)
+
+    def test_matches_with_type_mask(self) -> None:
+        type_mask = np.ones(self.ntypes, dtype=np.int64)
+        type_mask[2] = 0
+        ref_bias, ref_std = compute_stats_from_redu(
+            self.output_redu, self.natoms * type_mask.reshape(1, -1)
+        )
+        bias, std = self._accumulate().solve(type_mask=type_mask)
+        # the excluded type is left at the numerical zero of the min-norm solution
+        np.testing.assert_allclose(bias, ref_bias, rtol=1e-9, atol=1e-9)
+        np.testing.assert_allclose(std, ref_std, rtol=1e-9)
+
+    def test_repeated_compression_is_exact(self) -> None:
+        ref_bias, ref_std = compute_stats_from_redu(self.output_redu, self.natoms)
+        acc = ReduStatAccumulator(self.ntypes, self.ndim, [self.ndim])
+        acc._compress_every = 16
+        for start in range(self.natoms.shape[0]):
+            acc.add(self.output_redu[start : start + 1], self.natoms[start : start + 1])
+        bias, std = acc.solve()
+        np.testing.assert_allclose(bias, ref_bias, rtol=1e-9)
+        np.testing.assert_allclose(std, ref_std, rtol=1e-9)
+
+    def test_counts_every_frame(self) -> None:
+        acc = self._accumulate()
+        self.assertEqual(acc.nframes, self.natoms.shape[0])
+        np.testing.assert_array_equal(acc.natoms_total, self.natoms.sum(axis=0))
+
+    def test_output_shape_is_preserved(self) -> None:
+        rng = np.random.default_rng(0)
+        acc = ReduStatAccumulator(self.ntypes, 6, [2, 3])
+        acc.add(rng.random((10, 2, 3)), rng.integers(1, 5, (10, self.ntypes)))
+        bias, std = acc.solve()
+        self.assertEqual(bias.shape, (self.ntypes, 2, 3))
+        self.assertEqual(std.shape, (2, 3))
+
+    def test_empty_accumulator_is_rejected(self) -> None:
+        with self.assertRaises(ValueError):
+            ReduStatAccumulator(self.ntypes, self.ndim).solve()
+
+
+class TestReduStatScanner(unittest.TestCase):
+    def test_scan_is_performed_once_per_request(self) -> None:
+        calls = []
+
+        def scan_fn(ntypes, keys, intensive):
+            calls.append((ntypes, keys, intensive))
+            return unittest.mock.Mock(natoms_total=np.array([1, 0, 2]))
+
+        scanner = ReduStatScanner(scan_fn)
+        scanner.scan(3, ["energy"])
+        scanner.scan(3, ["energy"])
+        scanner.scan(3, ["energy"], intensive=True)
+        self.assertEqual(len(calls), 2)
+
+    def test_natoms_total_reuses_a_previous_scan(self) -> None:
+        counts = np.array([1, 0, 2])
+        scanner = ReduStatScanner(
+            lambda ntypes, keys, intensive: unittest.mock.Mock(natoms_total=counts)
+        )
+        scanner.scan(3, ["energy"])
+        np.testing.assert_array_equal(scanner.natoms_total(3), counts)

From a9907f7cd4b816e5da661cf8bf0f736d70c6334d Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 9 Sep 2026 10:12:40 +0800
Subject: [PATCH 2/5] fix: route data_stat_full in bridging and type-check the
 scanner handoff

Two CI failures on the previous commit.

`test_routing_covers_the_argcheck_schema` requires every model key to have an
explicit routing decision in `deepmd.utils.bridging`. `data_stat_full`
configures the learned model's statistics exactly like `data_stat_nbatch` and
`data_stat_protect`, so it joins them in `_LEARNED_CHILD_KEYS`.

`test_update_mode_replaces_orphaned_output_pair` passes a `Mock` sampler, which
answers `redu_stat_scanner` with an auto-created `Mock`; the bare `getattr` then
treated it as a scanner. Add `get_redu_stat_scanner`, which type-checks the
attribute, and use it everywhere the scanner is read or forwarded.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 .../dpmodel/atomic_model/base_atomic_model.py   |  9 +++++----
 deepmd/pd/model/atomic_model/dp_atomic_model.py |  7 ++++---
 deepmd/pd/utils/stat.py                         |  3 ++-
 .../pt/model/atomic_model/base_atomic_model.py  |  9 +++++----
 deepmd/pt/utils/stat.py                         |  3 ++-
 deepmd/utils/bridging.py                        |  1 +
 deepmd/utils/out_stat.py                        | 11 +++++++++++
 source/tests/common/test_out_stat.py            | 17 +++++++++++++++++
 8 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py
index 6587fe41ae..24bc56074d 100644
--- a/deepmd/dpmodel/atomic_model/base_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py
@@ -41,6 +41,9 @@
     map_atom_exclude_types,
     map_pair_exclude_types,
 )
+from deepmd.utils.out_stat import (
+    get_redu_stat_scanner,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -114,7 +117,7 @@ def _collect_and_set_observed_type(
         else:
             observed = _restore_observed_type_from_file(stat_file_path)
             if observed is None:
-                scanner = getattr(sampled_func, "redu_stat_scanner", None)
+                scanner = get_redu_stat_scanner(sampled_func)
                 if scanner is not None:
                     observed = observed_types_from_counts(
                         scanner.natoms_total(len(self.type_map)), self.type_map
@@ -775,9 +778,7 @@ def wrapped_sampler() -> list[dict]:
 
         # the full-data scanner, when the trainer attached one, is part of the
         # sampler contract and must survive wrapping
-        wrapped_sampler.redu_stat_scanner = getattr(
-            sampled_func, "redu_stat_scanner", None
-        )
+        wrapped_sampler.redu_stat_scanner = get_redu_stat_scanner(sampled_func)
         return wrapped_sampler
 
     def change_out_bias(
diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py
index 2dca9ea65f..ecde07ad69 100644
--- a/deepmd/pd/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py
@@ -20,6 +20,9 @@
 from deepmd.pd.model.task.base_fitting import (
     BaseFitting,
 )
+from deepmd.utils.out_stat import (
+    get_redu_stat_scanner,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -411,9 +414,7 @@ def wrapped_sampler() -> list[dict]:
 
         # the full-data scanner, when the trainer attached one, is part of the
         # sampler contract and must survive wrapping
-        wrapped_sampler.redu_stat_scanner = getattr(
-            sampled_func, "redu_stat_scanner", None
-        )
+        wrapped_sampler.redu_stat_scanner = get_redu_stat_scanner(sampled_func)
 
         self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path)
         self.compute_fitting_input_stat(wrapped_sampler, stat_file_path)
diff --git a/deepmd/pd/utils/stat.py b/deepmd/pd/utils/stat.py
index 1af3fd0933..27de6f7707 100644
--- a/deepmd/pd/utils/stat.py
+++ b/deepmd/pd/utils/stat.py
@@ -36,6 +36,7 @@
     compute_stats_do_not_distinguish_types,
     compute_stats_from_atomic,
     compute_stats_from_redu,
+    get_redu_stat_scanner,
 )
 from deepmd.utils.path import (
     DPPath,
@@ -470,7 +471,7 @@ def compute_output_stats(
             stats_distinguish_types,
             intensive,
             model_pred_g,
-            getattr(merged, "redu_stat_scanner", None),
+            get_redu_stat_scanner(merged),
         )
         bias_atom_a, std_atom_a = _compute_output_stats_atomic(
             sampled,
diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py
index 3da0e16425..d4f4b4efc7 100644
--- a/deepmd/pt/model/atomic_model/base_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/base_atomic_model.py
@@ -40,6 +40,9 @@
     map_atom_exclude_types,
     map_pair_exclude_types,
 )
+from deepmd.utils.out_stat import (
+    get_redu_stat_scanner,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -127,7 +130,7 @@ def _collect_and_set_observed_type(
         else:
             observed = _restore_observed_type_from_file(stat_file_path)
             if observed is None:
-                scanner = getattr(sampled_func, "redu_stat_scanner", None)
+                scanner = get_redu_stat_scanner(sampled_func)
                 if scanner is not None:
                     observed = observed_types_from_counts(
                         scanner.natoms_total(len(self.type_map)), self.type_map
@@ -262,9 +265,7 @@ def wrapped_sampler() -> list[dict]:
 
         # the full-data scanner, when the trainer attached one, is part of the
         # sampler contract and must survive wrapping
-        wrapped_sampler.redu_stat_scanner = getattr(
-            sampled_func, "redu_stat_scanner", None
-        )
+        wrapped_sampler.redu_stat_scanner = get_redu_stat_scanner(sampled_func)
         return wrapped_sampler
 
     def reinit_atom_exclude(
diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
index fdf7f87f09..d97f9b4deb 100644
--- a/deepmd/pt/utils/stat.py
+++ b/deepmd/pt/utils/stat.py
@@ -35,6 +35,7 @@
     compute_stats_do_not_distinguish_types,
     compute_stats_from_atomic,
     compute_stats_from_redu,
+    get_redu_stat_scanner,
 )
 from deepmd.utils.path import (
     DPPath,
@@ -672,7 +673,7 @@ def compute_output_stats(
             stats_distinguish_types,
             intensive,
             model_pred_g,
-            getattr(merged, "redu_stat_scanner", None),
+            get_redu_stat_scanner(merged),
         )
         bias_atom_a, std_atom_a = _compute_output_stats_atomic(
             sampled,
diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py
index 88467d66ab..fbdc3c103d 100644
--- a/deepmd/utils/bridging.py
+++ b/deepmd/utils/bridging.py
@@ -104,6 +104,7 @@ def _is_dpa4_family(sub: dict) -> bool:
     "enable_tf32",
     "data_stat_nbatch",
     "data_stat_protect",
+    "data_stat_full",
     "data_bias_nsample",
     "use_srtab",
     "smin_alpha",
diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py
index d62d681456..29006f2580 100644
--- a/deepmd/utils/out_stat.py
+++ b/deepmd/utils/out_stat.py
@@ -407,3 +407,14 @@ def natoms_total(self, ntypes: int) -> np.ndarray:
         if self._cache:
             return next(iter(self._cache.values())).natoms_total
         return self.scan(ntypes, ()).natoms_total
+
+
+def get_redu_stat_scanner(sampler: object) -> ReduStatScanner | None:
+    """Return the full-data scanner a statistics sampler carries, if any.
+
+    The scanner rides on the sampler as an attribute so that it reaches the
+    statistics consumers without a new argument on every atomic model. The type
+    check keeps that loose contract from picking up an unrelated attribute.
+    """
+    scanner = getattr(sampler, "redu_stat_scanner", None)
+    return scanner if isinstance(scanner, ReduStatScanner) else None
diff --git a/source/tests/common/test_out_stat.py b/source/tests/common/test_out_stat.py
index f6d2b3ab03..02db2b217d 100644
--- a/source/tests/common/test_out_stat.py
+++ b/source/tests/common/test_out_stat.py
@@ -7,6 +7,7 @@
 from deepmd.utils.out_stat import (
     ReduStatAccumulator,
     ReduStatScanner,
+    get_redu_stat_scanner,
     compute_stats_do_not_distinguish_types,
     compute_stats_from_atomic,
     compute_stats_from_redu,
@@ -317,3 +318,19 @@ def test_natoms_total_reuses_a_previous_scan(self) -> None:
         )
         scanner.scan(3, ["energy"])
         np.testing.assert_array_equal(scanner.natoms_total(3), counts)
+
+
+class TestGetReduStatScanner(unittest.TestCase):
+    def test_returns_an_attached_scanner(self) -> None:
+        def sampler() -> list:
+            return []
+
+        scanner = ReduStatScanner(lambda ntypes, keys, intensive: None)
+        sampler.redu_stat_scanner = scanner
+        self.assertIs(get_redu_stat_scanner(sampler), scanner)
+
+    def test_ignores_a_non_scanner_attribute(self) -> None:
+        # a Mock sampler answers every attribute; only a real scanner counts
+        self.assertIsNone(get_redu_stat_scanner(unittest.mock.Mock()))
+        self.assertIsNone(get_redu_stat_scanner(lambda: []))
+        self.assertIsNone(get_redu_stat_scanner([{}]))

From 36b955a4bb372845cadd2cc0da38907a2afd8bcf Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 9 Sep 2026 02:14:47 +0000
Subject: [PATCH 3/5] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/common/test_out_stat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/source/tests/common/test_out_stat.py b/source/tests/common/test_out_stat.py
index 02db2b217d..854e998f54 100644
--- a/source/tests/common/test_out_stat.py
+++ b/source/tests/common/test_out_stat.py
@@ -7,10 +7,10 @@
 from deepmd.utils.out_stat import (
     ReduStatAccumulator,
     ReduStatScanner,
-    get_redu_stat_scanner,
     compute_stats_do_not_distinguish_types,
     compute_stats_from_atomic,
     compute_stats_from_redu,
+    get_redu_stat_scanner,
 )
 
 

From 7af40477611e48c8a4de2514506ce074db27f475 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 9 Sep 2026 14:10:32 +0800
Subject: [PATCH 4/5] fix: honour data_stat_full over a sampled cache, forward
 intensive in pd

Addresses review feedback on #6011.

`compute_output_stats` returned early on any cached bias, so setting
`data_stat_full: true` against an existing stat file silently reused the sampled
values the flag exists to replace. The cache now records how its output
statistics were produced: a full-scan run rejects a sampled cache and recomputes,
a sampled recomputation withdraws an earlier full-scan claim, and a cache that
never held full-scan values keeps the legacy layout. A read-only cache cannot be
recomputed, so it warns instead.

Resolving the scanner in `compute_output_stats` also puts the delta-bias and
type-blind fallback where the information lives, ahead of the cache lookup.

Drive-by, flagged in review: the pd sampled path dropped `intensive` when calling
`compute_stats_from_redu`, which the pt path and the new scan path both honour.

The `type_mask` built from `sampled[0]` was also raised; `atom_exclude_types` is
written to every sample from one model-level source, so it is uniform by
construction. Noted in a comment rather than changed.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 deepmd/pd/utils/stat.py                | 63 +++++++++++++++++----
 deepmd/pt/utils/stat.py                | 60 ++++++++++++++++----
 deepmd/utils/stat_file.py              | 46 ++++++++++++++++
 source/tests/pt/test_stat_file_mode.py | 76 ++++++++++++++++++++++++++
 4 files changed, 223 insertions(+), 22 deletions(-)

diff --git a/deepmd/pd/utils/stat.py b/deepmd/pd/utils/stat.py
index 27de6f7707..4bd3cc22ff 100644
--- a/deepmd/pd/utils/stat.py
+++ b/deepmd/pd/utils/stat.py
@@ -41,6 +41,10 @@
 from deepmd.utils.path import (
     DPPath,
 )
+from deepmd.utils.stat_file import (
+    load_output_stat_full_scan,
+    save_output_stat_full_scan,
+)
 
 log = logging.getLogger(__name__)
 
@@ -153,6 +157,11 @@ def scan_redu_stats(
     -------
     ReduScanResult
         The accumulators, the per-type atom counts and the frame count.
+
+    Notes
+    -----
+    Statistics initialization runs on the chief process only, so one scan of
+    the training set is performed per run, not per rank.
     """
     stats: dict[str, ReduStatAccumulator] = {}
     natoms_total = np.zeros(ntypes, dtype=np.int64)
@@ -384,8 +393,39 @@ def compute_output_stats(
     intensive : bool, optional
         Whether the fitting target is intensive.
     """
+    # a full scan cannot replace the sampled path for delta bias or type-blind
+    # statistics, so resolve it before the cache is consulted
+    redu_scanner = get_redu_stat_scanner(merged)
+    if redu_scanner is not None and (
+        model_forward is not None or not stats_distinguish_types
+    ):
+        log.warning(
+            "Falling back to sampled output statistics: a full scan supports "
+            "neither delta bias nor stats_distinguish_types=False."
+        )
+        redu_scanner = None
+
     # try to restore the bias from stat file
     bias_atom_e, std_atom_e = _restore_from_file(stat_file_path, keys)
+    if (
+        bias_atom_e is not None
+        and redu_scanner is not None
+        and not load_output_stat_full_scan(stat_file_path)
+    ):
+        # the cache holds sampled values, which is what data_stat_full replaces
+        if getattr(stat_file_path, "mode", None) == "r":
+            log.warning(
+                "`data_stat_full` is set, but the read-only statistics cache "
+                "holds output statistics estimated from sampled batches; they "
+                "are used as they are."
+            )
+        else:
+            log.info(
+                "Recomputing output statistics: the cache holds values "
+                "estimated from sampled batches, which `data_stat_full` "
+                "replaces."
+            )
+            bias_atom_e, std_atom_e = None, None
 
     # failed to restore the bias from stat file. compute
     if bias_atom_e is None:
@@ -471,7 +511,7 @@ def compute_output_stats(
             stats_distinguish_types,
             intensive,
             model_pred_g,
-            get_redu_stat_scanner(merged),
+            redu_scanner,
         )
         bias_atom_a, std_atom_a = _compute_output_stats_atomic(
             sampled,
@@ -502,6 +542,7 @@ def compute_output_stats(
 
         if stat_file_path is not None:
             _save_to_file(stat_file_path, bias_atom_e, std_atom_e)
+            save_output_stat_full_scan(stat_file_path, redu_scanner is not None)
 
     bias_atom_e = {kk: to_paddle_tensor(vv) for kk, vv in bias_atom_e.items()}
     std_atom_e = {kk: to_paddle_tensor(vv) for kk, vv in std_atom_e.items()}
@@ -524,7 +565,8 @@ def _compute_output_stats_global(
 
     When *redu_scanner* is given, the bias and std come from a scan of every
     training frame instead of the sampled batches, which keeps rare elements
-    from being under-represented in the regression.
+    from being under-represented in the regression. The caller decides whether
+    a scan is admissible; it is ignored for delta bias and type-blind statistics.
     """
     # return directly if no global samples
     if global_sampled_idx is None or all(
@@ -585,15 +627,13 @@ def _compute_output_stats_global(
             if kk in merged_output
         }
 
-    scan = None
-    if redu_scanner is not None:
-        if model_pred is not None or not stats_distinguish_types:
-            log.warning(
-                "Falling back to sampled output statistics: a full scan supports "
-                "neither delta bias nor stats_distinguish_types=False."
-            )
-        else:
-            scan = redu_scanner.scan(ntypes, keys, intensive)
+    scan = (
+        redu_scanner.scan(ntypes, keys, intensive)
+        if redu_scanner is not None and model_pred is None and stats_distinguish_types
+        else None
+    )
+    # one model-level source writes atom_exclude_types onto every sample, so
+    # the first one carries the mask the whole scan needs
     type_mask = (
         to_numpy_array(
             AtomExcludeMask(ntypes, sampled[0]["atom_exclude_types"]).get_type_mask()
@@ -629,6 +669,7 @@ def _compute_output_stats_global(
                     merged_natoms[kk],
                     assigned_bias=assigned_atom_ener[kk],
                     rcond=rcond,
+                    intensive=intensive,
                 )
         else:
             # this key does not have global labels, skip it.
diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
index d97f9b4deb..33ca212918 100644
--- a/deepmd/pt/utils/stat.py
+++ b/deepmd/pt/utils/stat.py
@@ -41,8 +41,10 @@
     DPPath,
 )
 from deepmd.utils.stat_file import (
+    load_output_stat_full_scan,
     load_paired_items,
     replace_paired_items,
+    save_output_stat_full_scan,
 )
 
 log = logging.getLogger(__name__)
@@ -286,6 +288,11 @@ def scan_redu_stats(
     -------
     ReduScanResult
         The accumulators, the per-type atom counts and the frame count.
+
+    Notes
+    -----
+    Statistics initialization runs on the chief process only, so one scan of
+    the training set is performed per run, not per rank.
     """
     stats: dict[str, ReduStatAccumulator] = {}
     natoms_total = np.zeros(ntypes, dtype=np.int64)
@@ -574,8 +581,39 @@ def compute_output_stats(
     assert isinstance(keys, list)
     requested_keys = list(keys)
 
+    # a full scan cannot replace the sampled path for delta bias or type-blind
+    # statistics, so resolve it before the cache is consulted
+    redu_scanner = get_redu_stat_scanner(merged)
+    if redu_scanner is not None and (
+        model_forward is not None or not stats_distinguish_types
+    ):
+        log.warning(
+            "Falling back to sampled output statistics: a full scan supports "
+            "neither delta bias nor stats_distinguish_types=False."
+        )
+        redu_scanner = None
+
     # try to restore the bias from stat file
     bias_atom_e, std_atom_e = _restore_from_file(stat_file_path, keys)
+    if (
+        bias_atom_e is not None
+        and redu_scanner is not None
+        and not load_output_stat_full_scan(stat_file_path)
+    ):
+        # the cache holds sampled values, which is what data_stat_full replaces
+        if getattr(stat_file_path, "mode", None) == "r":
+            log.warning(
+                "`data_stat_full` is set, but the read-only statistics cache "
+                "holds output statistics estimated from sampled batches; they "
+                "are used as they are."
+            )
+        else:
+            log.info(
+                "Recomputing output statistics: the cache holds values "
+                "estimated from sampled batches, which `data_stat_full` "
+                "replaces."
+            )
+            bias_atom_e, std_atom_e = None, None
 
     # failed to restore the bias from stat file. compute
     if bias_atom_e is None:
@@ -673,7 +711,7 @@ def compute_output_stats(
             stats_distinguish_types,
             intensive,
             model_pred_g,
-            get_redu_stat_scanner(merged),
+            redu_scanner,
         )
         bias_atom_a, std_atom_a = _compute_output_stats_atomic(
             sampled,
@@ -709,6 +747,7 @@ def compute_output_stats(
                 bias_atom_e,
                 std_atom_e,
             )
+            save_output_stat_full_scan(stat_file_path, redu_scanner is not None)
 
     bias_atom_e = {kk: to_torch_tensor(vv) for kk, vv in bias_atom_e.items()}
     std_atom_e = {kk: to_torch_tensor(vv) for kk, vv in std_atom_e.items()}
@@ -731,7 +770,8 @@ def _compute_output_stats_global(
 
     When *redu_scanner* is given, the bias and std come from a scan of every
     training frame instead of the sampled batches, which keeps rare elements
-    from being under-represented in the regression.
+    from being under-represented in the regression. The caller decides whether
+    a scan is admissible; it is ignored for delta bias and type-blind statistics.
     """
     # return directly if no global samples
     if global_sampled_idx is None or all(
@@ -792,15 +832,13 @@ def _compute_output_stats_global(
             if kk in merged_output
         }
 
-    scan = None
-    if redu_scanner is not None:
-        if model_pred is not None or not stats_distinguish_types:
-            log.warning(
-                "Falling back to sampled output statistics: a full scan supports "
-                "neither delta bias nor stats_distinguish_types=False."
-            )
-        else:
-            scan = redu_scanner.scan(ntypes, keys, intensive)
+    scan = (
+        redu_scanner.scan(ntypes, keys, intensive)
+        if redu_scanner is not None and model_pred is None and stats_distinguish_types
+        else None
+    )
+    # one model-level source writes atom_exclude_types onto every sample, so
+    # the first one carries the mask the whole scan needs
     type_mask = (
         to_numpy_array(
             AtomExcludeMask(ntypes, sampled[0]["atom_exclude_types"]).get_type_mask()
diff --git a/deepmd/utils/stat_file.py b/deepmd/utils/stat_file.py
index 98e02f72e7..ee3b4e750e 100644
--- a/deepmd/utils/stat_file.py
+++ b/deepmd/utils/stat_file.py
@@ -202,6 +202,52 @@ def open_stat_file(
         owner.close()
 
 
+_FULL_SCAN_ITEM = "output_stat_full_scan"
+
+
+def load_output_stat_full_scan(path: DPPath | None) -> bool:
+    """Report whether the cached output statistics came from a full data scan.
+
+    A cache written before this flag existed reports ``False``: its values were
+    estimated from sampled batches.
+
+    Parameters
+    ----------
+    path
+        Statistics-cache root used by the current consumer.
+
+    Returns
+    -------
+    bool
+        Whether the cached output statistics scanned every frame.
+    """
+    if path is None or not (path / _FULL_SCAN_ITEM).is_file():
+        return False
+    return bool(np.asarray((path / _FULL_SCAN_ITEM).load_numpy()).item())
+
+
+def save_output_stat_full_scan(path: DPPath | None, full_scan: bool) -> None:
+    """Record how the output statistics now in the cache were produced.
+
+    A cache that never held full-scan statistics keeps the legacy layout: the
+    absence of the item already means that its values were sampled. The item is
+    written only to claim a full scan, or to withdraw a claim that a sampled
+    recomputation has just invalidated.
+
+    Parameters
+    ----------
+    path
+        Writable statistics-cache root.
+    full_scan
+        Whether the stored statistics scanned every frame.
+    """
+    if path is None or (not full_scan and not load_output_stat_full_scan(path)):
+        return
+    path.mkdir(exist_ok=True, parents=True)
+    # a one-element array: the cache readers slice what they load
+    (path / _FULL_SCAN_ITEM).save_numpy(np.array([bool(full_scan)]))
+
+
 def load_required_items(
     path: DPPath | None,
     names: Sequence[str],
diff --git a/source/tests/pt/test_stat_file_mode.py b/source/tests/pt/test_stat_file_mode.py
index 706ce1381b..4d63312162 100644
--- a/source/tests/pt/test_stat_file_mode.py
+++ b/source/tests/pt/test_stat_file_mode.py
@@ -37,8 +37,14 @@
 from deepmd.utils.argcheck import (
     normalize,
 )
+from deepmd.utils.out_stat import (
+    ReduScanResult,
+    ReduStatAccumulator,
+    ReduStatScanner,
+)
 from deepmd.utils.stat_file import (
     StatFileSpec,
+    load_output_stat_full_scan,
     open_stat_file,
 )
 
@@ -122,6 +128,76 @@ def _energy_stat_sample() -> list[dict[str, Any]]:
     ]
 
 
+def _energy_scan_sampler(bias: float) -> Mock:
+    """A stat sampler carrying a full-data scanner that fits a constant bias."""
+    sampler = Mock(return_value=_energy_stat_sample())
+    natoms = np.array([[2, 0], [0, 2]], dtype=np.float64)
+
+    def scan(ntypes: int, keys: tuple, intensive: bool) -> ReduScanResult:
+        accumulator = ReduStatAccumulator(ntypes, 1, [1], intensive=intensive)
+        accumulator.add(natoms.sum(axis=1, keepdims=True) * bias, natoms)
+        return ReduScanResult(
+            stats={"energy": accumulator},
+            natoms_total=natoms.sum(axis=0).astype(np.int64),
+            nframes=natoms.shape[0],
+        )
+
+    sampler.redu_stat_scanner = ReduStatScanner(scan)
+    return sampler
+
+
+def test_full_scan_replaces_a_sampled_cache(tmp_path: Path) -> None:
+    stat_file = tmp_path / "stat.hdf5"
+    sampled = _compute_energy_stats(stat_file, Mock(return_value=_energy_stat_sample()))
+    # a cache that never held full-scan values keeps the legacy layout
+    assert not _full_scan_claim(stat_file)
+    with h5py.File(stat_file, "r") as file:
+        assert set(file) == {"bias_atom_energy", "std_atom_energy"}
+
+    # those values were estimated from batches, which is what the flag exists
+    # to replace, so the scan must win over the cache
+    scanned = _compute_energy_stats(stat_file, _energy_scan_sampler(7.0))
+    assert _full_scan_claim(stat_file)
+    np.testing.assert_allclose(scanned, 7.0)
+    assert not np.allclose(sampled, 7.0)
+
+    # a second full-scan run is a cache hit and must neither sample nor rescan
+    rescanner = _energy_scan_sampler(9.0)
+    np.testing.assert_allclose(_compute_energy_stats(stat_file, rescanner), 7.0)
+    rescanner.assert_not_called()
+
+
+def _full_scan_claim(stat_file: Path) -> bool:
+    with open_stat_file(StatFileSpec(str(stat_file), "read")) as stat_path:
+        return load_output_stat_full_scan(stat_path)
+
+
+def _compute_energy_stats(stat_file: Path, sampler: Mock) -> np.ndarray:
+    with open_stat_file(StatFileSpec(str(stat_file))) as stat_path:
+        assert stat_path is not None
+        bias, _ = compute_output_stats(
+            sampler,
+            ntypes=2,
+            keys=["energy"],
+            stat_file_path=stat_path,
+        )
+    return bias["energy"].cpu().numpy()
+
+
+def test_sampled_recompute_withdraws_the_full_scan_claim(tmp_path: Path) -> None:
+    stat_file = tmp_path / "stat.hdf5"
+    _compute_energy_stats(stat_file, _energy_scan_sampler(7.0))
+    assert _full_scan_claim(stat_file)
+
+    # an incomplete pair forces a recomputation, here from sampled batches
+    with h5py.File(stat_file, "a") as file:
+        del file["std_atom_energy"]
+    sampled = _compute_energy_stats(stat_file, Mock(return_value=_energy_stat_sample()))
+
+    assert not _full_scan_claim(stat_file)
+    assert not np.allclose(sampled, 7.0)
+
+
 def test_default_stat_file_mode_remains_writable(tmp_path: Path) -> None:
     stat_file = tmp_path / "stat.hdf5"
     with open_stat_file(StatFileSpec(str(stat_file))) as stat_path:

From f1e5614063abee26a787f812ff03ee179b9f0206 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Wed, 9 Sep 2026 15:52:54 +0800
Subject: [PATCH 5/5] fix: withdraw the full-scan claim before the values it
 describes are replaced

Raised in review on #6011. The claim was updated after `_save_to_file`, so an
interruption between the two left sampled values marked as full-scan ones, and
the next `data_stat_full` run would accept that cache and skip the scan it was
asked for.

Clear the claim first. An interruption now leaves the cache looking sampled,
which costs a rescan instead of silently returning the wrong statistics.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 deepmd/pd/utils/stat.py                |  3 +++
 deepmd/pt/utils/stat.py                |  3 +++
 source/tests/pt/test_stat_file_mode.py | 21 +++++++++++++++++++++
 3 files changed, 27 insertions(+)

diff --git a/deepmd/pd/utils/stat.py b/deepmd/pd/utils/stat.py
index 4bd3cc22ff..654a28c85d 100644
--- a/deepmd/pd/utils/stat.py
+++ b/deepmd/pd/utils/stat.py
@@ -541,6 +541,9 @@ def compute_output_stats(
                 raise RuntimeError("Fail to compute stat.")
 
         if stat_file_path is not None:
+            # withdraw any standing claim before the values it describes are
+            # replaced, so an interruption leaves the cache looking sampled
+            save_output_stat_full_scan(stat_file_path, False)
             _save_to_file(stat_file_path, bias_atom_e, std_atom_e)
             save_output_stat_full_scan(stat_file_path, redu_scanner is not None)
 
diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
index 33ca212918..ea5dbf1ebe 100644
--- a/deepmd/pt/utils/stat.py
+++ b/deepmd/pt/utils/stat.py
@@ -741,6 +741,9 @@ def compute_output_stats(
                 raise RuntimeError("Fail to compute stat.")
 
         if stat_file_path is not None:
+            # withdraw any standing claim before the values it describes are
+            # replaced, so an interruption leaves the cache looking sampled
+            save_output_stat_full_scan(stat_file_path, False)
             _save_to_file(
                 stat_file_path,
                 requested_keys,
diff --git a/source/tests/pt/test_stat_file_mode.py b/source/tests/pt/test_stat_file_mode.py
index 4d63312162..5645fb9361 100644
--- a/source/tests/pt/test_stat_file_mode.py
+++ b/source/tests/pt/test_stat_file_mode.py
@@ -198,6 +198,27 @@ def test_sampled_recompute_withdraws_the_full_scan_claim(tmp_path: Path) -> None
     assert not np.allclose(sampled, 7.0)
 
 
+def test_interrupted_replacement_drops_the_full_scan_claim(tmp_path: Path) -> None:
+    stat_file = tmp_path / "stat.hdf5"
+    _compute_energy_stats(stat_file, _energy_scan_sampler(7.0))
+    assert _full_scan_claim(stat_file)
+
+    with h5py.File(stat_file, "a") as file:
+        del file["std_atom_energy"]
+    with (
+        patch(
+            "deepmd.pt.utils.stat._save_to_file",
+            side_effect=RuntimeError("interrupted"),
+        ),
+        pytest.raises(RuntimeError, match="interrupted"),
+    ):
+        _compute_energy_stats(stat_file, Mock(return_value=_energy_stat_sample()))
+
+    # the claim describes values that were about to be overwritten, so it must
+    # not outlive them: a lost claim only costs a rescan, a stale one is wrong
+    assert not _full_scan_claim(stat_file)
+
+
 def test_default_stat_file_mode_remains_writable(tmp_path: Path) -> None:
     stat_file = tmp_path / "stat.hdf5"
     with open_stat_file(StatFileSpec(str(stat_file))) as stat_path: