From d8a4b67fb48b563816f9914fd7101fa6f5986a80 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Fri, 9 Oct 2026 23:52:00 +0800
Subject: [PATCH 1/6] feat(train): shared MetricAccumulator for disp_avg
 intervals

Generalize training display aggregation into a weighted MetricAccumulator
used by both PT and pt_expt. Training steps keep unit weight; validation
batches use atom weights. NaN optional-label metrics are excluded from
per-metric means instead of poisoning the interval. Distributed runs
all-reduce packed sums, weights, and observation counts only at display
boundaries.

Refs #5819 #5755
---
 deepmd/dpmodel/train/__init__.py              |   2 +
 deepmd/dpmodel/train/metrics.py               | 176 ++++++++++++++----
 deepmd/dpmodel/train/trainer.py               |  17 +-
 deepmd/pt/train/training.py                   |  48 +++--
 deepmd/pt_expt/train/metrics.py               |  57 ++++++
 deepmd/pt_expt/train/training.py              |  48 +++--
 deepmd/utils/argcheck.py                      |   8 +-
 .../common/dpmodel/test_train_metrics.py      |  44 ++++-
 .../tests/pt_expt/train/test_metric_reduce.py |  20 ++
 9 files changed, 350 insertions(+), 70 deletions(-)
 create mode 100644 deepmd/pt_expt/train/metrics.py
 create mode 100644 source/tests/pt_expt/train/test_metric_reduce.py

diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py
index 1463dad928..8b880bff56 100644
--- a/deepmd/dpmodel/train/__init__.py
+++ b/deepmd/dpmodel/train/__init__.py
@@ -17,6 +17,7 @@
     TrainEntrypointOptions,
 )
 from .metrics import (
+    MetricAccumulator,
     TrainingMetricAccumulator,
 )
 from .schedule import (
@@ -50,6 +51,7 @@
     "CheckpointStore",
     "DisplayInterval",
     "LearningCurveWriter",
+    "MetricAccumulator",
     "RankContext",
     "ShardingPolicy",
     "StepSchedule",
diff --git a/deepmd/dpmodel/train/metrics.py b/deepmd/dpmodel/train/metrics.py
index 3ea0f405b8..ab32b1f574 100644
--- a/deepmd/dpmodel/train/metrics.py
+++ b/deepmd/dpmodel/train/metrics.py
@@ -1,23 +1,54 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""Backend-independent accumulation of training display metrics."""
+"""Backend-independent accumulation of training display metrics.
+
+Two weighting policies share one accumulator:
+
+* Training steps use unit weight, so the reported value is the arithmetic
+  mean of the detached per-step metrics.
+* Validation batches use an atom weight (real atoms in the batch), so the
+  reported value is the atom-weighted mean of per-atom metrics.
+
+Optional labels emit NaN through ``display_if_exist``. Those observations are
+excluded from that metric's sum and weight so a single unlabeled batch cannot
+poison an otherwise finite interval. Distributed runs keep local sums and
+weights on device and reduce them only at a display boundary.
+"""
+
+from __future__ import annotations
 
 from collections.abc import (
     Mapping,
     Sequence,
 )
+from math import (
+    isnan,
+)
 from typing import (
     Any,
 )
 
 
-class TrainingMetricAccumulator:
-    """Average per-step metrics independently for each training task.
+def _is_nan(value: Any) -> bool:
+    """Return whether a scalar host or array value is NaN."""
+    if value is None:
+        return True
+    try:
+        return bool(isnan(float(value)))
+    except (TypeError, ValueError):
+        return False
+
 
-    Metric names are declared before training so an unsampled task has the
-    same columns as a sampled task. Each update has equal weight, matching
-    the arithmetic mean of the reported per-step metrics rather than a
-    pooled error over atoms. Scalar arrays stay on their original device
-    until :meth:`average` converts them to Python floats for display.
+class MetricAccumulator:
+    """Accumulate weighted sums and counts of detached display metrics.
+
+    Metric names are declared before training so an unsampled task keeps the
+    same columns as a sampled task. Each call to :meth:`add` is one
+    observation: training uses ``weight=1``, validation passes the number of
+    real atoms. A NaN metric is skipped for that observation only.
+
+    Scalar arrays stay on their original device until :meth:`average`
+    converts them to Python floats for display. Out-of-place addition
+    preserves recorded values when a backend reuses an output buffer.
 
     Parameters
     ----------
@@ -26,44 +57,123 @@ class TrainingMetricAccumulator:
     """
 
     def __init__(self, metric_names: Mapping[str, Sequence[str]]) -> None:
-        self._totals: dict[str, dict[str, Any]] = {
-            task: dict.fromkeys(sorted(names), 0.0)
-            for task, names in metric_names.items()
+        self._names: dict[str, tuple[str, ...]] = {
+            task: tuple(sorted(names)) for task, names in metric_names.items()
+        }
+        self._sums: dict[str, dict[str, Any]] = {
+            task: dict.fromkeys(names, 0.0) for task, names in self._names.items()
         }
-        self._counts = dict.fromkeys(metric_names, 0)
+        self._weights: dict[str, dict[str, Any]] = {
+            task: dict.fromkeys(names, 0.0) for task, names in self._names.items()
+        }
+        self._observations = dict.fromkeys(self._names, 0)
+
+    @property
+    def task_keys(self) -> tuple[str, ...]:
+        """Return the declared task keys in insertion order."""
+        return tuple(self._names)
 
-    def add(self, task_key: str, metrics: Mapping[str, Any]) -> None:
-        """Record one optimizer step's detached scalar metrics.
+    def metric_names(self, task_key: str) -> tuple[str, ...]:
+        """Return the declared metric names for one task."""
+        return self._names[task_key]
+
+    def add(
+        self,
+        task_key: str,
+        metrics: Mapping[str, Any],
+        *,
+        weight: float = 1.0,
+    ) -> None:
+        """Record one observation of detached scalar metrics.
 
         Backends detach metrics from their differentiation graph before
-        calling this method. Missing metrics propagate NaN rather than being
-        treated as zero. Out-of-place addition preserves the recorded value
-        when a backend reuses an output buffer on its next step.
+        calling this method. Missing or NaN metrics are omitted from that
+        metric's weighted average rather than treated as zero. A non-positive
+        weight records the observation for task sampling but contributes to
+        no metric totals.
         """
-        totals = self._totals[task_key]
-        unknown = metrics.keys() - totals.keys()
+        names = self._names[task_key]
+        unknown = metrics.keys() - set(names)
         if unknown:
             raise ValueError(
                 f"Undeclared display metrics for task {task_key!r}: {sorted(unknown)}"
             )
-        for name, total in totals.items():
-            totals[name] = total + metrics.get(name, float("nan"))
-        self._counts[task_key] += 1
+        sums = self._sums[task_key]
+        weights = self._weights[task_key]
+        for name in names:
+            value = metrics.get(name, float("nan"))
+            if weight <= 0.0 or _is_nan(value):
+                continue
+            # Out-of-place update: backends may overwrite the tensor buffer
+            # that produced ``value`` on the next step.
+            sums[name] = sums[name] + value * weight
+            weights[name] = weights[name] + weight
+        self._observations[task_key] += 1
 
     def count(self, task_key: str) -> int:
-        """Return the number of recorded optimizer steps for one task."""
-        return self._counts[task_key]
+        """Return the number of recorded observations for one task.
+
+        This is the sampling count used to detect tasks that never ran in a
+        display interval. It is independent of per-metric weights.
+        """
+        return self._observations[task_key]
 
     def average(self, task_key: str) -> dict[str, float]:
-        """Return interval averages, or NaN for every unsampled metric."""
-        totals = self._totals[task_key]
-        count = self._counts[task_key]
-        if count == 0:
-            return dict.fromkeys(totals, float("nan"))
-        return {name: float(total / count) for name, total in totals.items()}
+        """Return weighted averages, or NaN for every metric with zero weight."""
+        sums = self._sums[task_key]
+        weights = self._weights[task_key]
+        result: dict[str, float] = {}
+        for name in self._names[task_key]:
+            weight = float(weights[name])
+            if weight == 0.0 or _is_nan(weight):
+                result[name] = float("nan")
+            else:
+                result[name] = float(sums[name] / weights[name])
+        return result
 
     def reset(self) -> None:
         """Clear interval values and counts while preserving metric names."""
-        for task, totals in self._totals.items():
-            self._totals[task] = dict.fromkeys(totals, 0.0)
-            self._counts[task] = 0
+        for task, names in self._names.items():
+            self._sums[task] = dict.fromkeys(names, 0.0)
+            self._weights[task] = dict.fromkeys(names, 0.0)
+            self._observations[task] = 0
+
+    def flat_state(self) -> list[Any]:
+        """Pack sums, weights, and observation counts for a collective reduce.
+
+        Layout is stable across ranks that share the same constructor
+        arguments: for each task, all metric sums, then all metric weights,
+        then the observation count.
+        """
+        values: list[Any] = []
+        for task, names in self._names.items():
+            sums = self._sums[task]
+            weights = self._weights[task]
+            values.extend(sums[name] for name in names)
+            values.extend(weights[name] for name in names)
+            values.append(float(self._observations[task]))
+        return values
+
+    def load_flat_state(self, values: Sequence[Any]) -> None:
+        """Replace local state with a previously packed, possibly reduced vector."""
+        expected = sum(2 * len(names) + 1 for names in self._names.values())
+        if len(values) != expected:
+            raise ValueError(
+                f"MetricAccumulator state length {len(values)} != expected {expected}"
+            )
+        index = 0
+        for task, names in self._names.items():
+            sums = self._sums[task]
+            weights = self._weights[task]
+            for name in names:
+                sums[name] = values[index]
+                index += 1
+            for name in names:
+                weights[name] = values[index]
+                index += 1
+            self._observations[task] = int(float(values[index]))
+            index += 1
+
+
+# Historical name kept for call sites and tests that predate the weighted API.
+TrainingMetricAccumulator = MetricAccumulator
diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py
index 531bc7d07c..1e22103ca9 100644
--- a/deepmd/dpmodel/train/trainer.py
+++ b/deepmd/dpmodel/train/trainer.py
@@ -56,6 +56,7 @@
 
 if TYPE_CHECKING:
     from .metrics import (
+        MetricAccumulator,
         TrainingMetricAccumulator,
     )
 
@@ -537,7 +538,7 @@ def __init__(
         trainer_config: TrainerConfig,
         *,
         rank_context: RankContext | None = None,
-        metric_accumulator: TrainingMetricAccumulator | None = None,
+        metric_accumulator: MetricAccumulator | None = None,
     ) -> None:
         self.trainer_config = trainer_config
         self.rank_context = rank_context or RankContext()
@@ -572,6 +573,10 @@ def run(self, tasks: TrainingTaskCollection) -> None:
                 display_step = step + 1
 
                 if self._should_display(display_step):
+                    # Collectives for averaged metrics run on every rank before
+                    # the chief-only logging path. Local windows stay on device
+                    # until this boundary.
+                    self.synchronize_metric_accumulator()
                     if self.rank_context.is_chief:
                         train_results, valid_results = self.collect_display_results(
                             tasks,
@@ -665,6 +670,16 @@ def collect_display_results(
             return train_results[active_task.key], valid_results[active_task.key]
         return train_results, valid_results
 
+    def synchronize_metric_accumulator(self) -> None:
+        """Reduce local metric windows across ranks at a display boundary.
+
+        The default implementation is a no-op for single-process runs. A
+        distributed backend overrides this with an all-reduce of the packed
+        sums, weights, and observation counts. The call is collective: every
+        rank that accumulates metrics must enter it together.
+        """
+        return None
+
     def on_train_begin(self, tasks: TrainingTaskCollection) -> None:
         """Hook called before the first optimizer step."""
         return None
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index 7a7210962e..b0200515b3 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -28,6 +28,7 @@
 from deepmd.dpmodel.train import (
     DEFAULT_TASK_KEY,
     CheckpointStore,
+    MetricAccumulator,
     ShardingPolicy,
     TrainingMetricAccumulator,
     TrainingTimer,
@@ -35,6 +36,9 @@
     change_model_out_bias,
     resolve_step_schedule,
 )
+from deepmd.pt_expt.train.metrics import (
+    all_reduce_metric_accumulator,
+)
 from deepmd.dpmodel.utils import (
     compute_total_numb_batch,
 )
@@ -1593,22 +1597,28 @@ def fake_model() -> dict:
                 display_step_id % self.disp_freq == 0 or display_step_id == 1
             ):
                 self.wrapper.eval()  # Will set to train mode before fininshing validation
+                if self.metric_accumulator is not None and self.world_size > 1:
+                    all_reduce_metric_accumulator(self.metric_accumulator)
 
                 def log_loss_valid(_task_key: str = "Default") -> dict:
-                    single_results = {}
-                    sum_natoms = 0
                     if not self.multi_task:
                         valid_numb_batch = self.valid_numb_batch
+                        loss_obj = self.loss
                     else:
                         valid_numb_batch = self.valid_numb_batch[_task_key]
+                        loss_obj = self.loss[_task_key]
+                    accumulator = MetricAccumulator(
+                        {_task_key: loss_obj.training_metric_names}
+                    )
+                    saw_batch = False
                     for ii in range(valid_numb_batch):
                         self.optimizer.zero_grad()
                         input_dict, label_dict, _ = self.get_data(
                             is_train=False, task_key=_task_key
                         )
                         if input_dict == {}:
-                            # no validation data
-                            return {}
+                            # no validation data for this slot
+                            break
                         # Validation runs the inner module, not the DDP
                         # wrapper. A DDP forward under grad mode arms the
                         # reducer for an all-reduce that this loop never
@@ -1624,20 +1634,26 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
                             label=label_dict,
                             task_key=_task_key,
                         )
-                        # more_loss.update({"rmse": math.sqrt(loss)})
-                        # The metrics are per-atom quantities, so each batch
-                        # weighs by the real atoms it holds summed over its
-                        # frames. Phantom atoms (atype < 0), which pad a
-                        # mixed-nloc batch, contribute to none of them.
+                        # Per-atom metrics weigh by real atoms across frames.
+                        # Phantom atoms (atype < 0) pad mixed-nloc batches.
                         natoms = int((input_dict["atype"] >= 0).sum())
-                        sum_natoms += natoms
-                        for k, v in more_loss.items():
-                            if "l2_" not in k:
-                                single_results[k] = (
-                                    single_results.get(k, 0.0) + v * natoms
+                        accumulator.add(
+                            _task_key,
+                            {
+                                name: (
+                                    value.detach()
+                                    if torch.is_tensor(value)
+                                    else value
                                 )
-                    results = {k: v / sum_natoms for k, v in single_results.items()}
-                    return results
+                                for name, value in more_loss.items()
+                                if "l2_" not in name
+                            },
+                            weight=float(natoms),
+                        )
+                        saw_batch = True
+                    if not saw_batch:
+                        return {}
+                    return accumulator.average(_task_key)
 
                 if not self.multi_task:
                     train_results = self._training_results(more_loss)
diff --git a/deepmd/pt_expt/train/metrics.py b/deepmd/pt_expt/train/metrics.py
new file mode 100644
index 0000000000..c78b7523b8
--- /dev/null
+++ b/deepmd/pt_expt/train/metrics.py
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Torch helpers for distributed metric-window reduction."""
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+import torch
+import torch.distributed as dist
+
+if TYPE_CHECKING:
+    from deepmd.dpmodel.train.metrics import (
+        MetricAccumulator,
+    )
+
+
+def all_reduce_metric_accumulator(
+    accumulator: MetricAccumulator,
+    *,
+    group: Any | None = None,
+) -> None:
+    """Sum local metric windows across ranks in place.
+
+    Call only at a display boundary. Host synchronization of the averaged
+    scalars still happens later in :meth:`MetricAccumulator.average`; this
+    routine only exchanges the packed sums, weights, and observation counts.
+    Single-process runs leave the accumulator unchanged.
+    """
+    if not dist.is_available() or not dist.is_initialized():
+        return
+    world_size = dist.get_world_size(group)
+    if world_size <= 1:
+        return
+
+    packed = accumulator.flat_state()
+    if not packed:
+        return
+
+    device = _state_device(packed)
+    tensor = torch.tensor(
+        [float(value) for value in packed],
+        dtype=torch.float64,
+        device=device,
+    )
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group)
+    accumulator.load_flat_state(tensor.detach().cpu().tolist())
+
+
+def _state_device(packed: list[Any]) -> torch.device:
+    """Prefer the device of the first tensor entry; otherwise use CPU."""
+    for value in packed:
+        if torch.is_tensor(value):
+            return value.device
+    return torch.device("cpu")
diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py
index 9130f5753b..d8fbfe5c85 100644
--- a/deepmd/pt_expt/train/training.py
+++ b/deepmd/pt_expt/train/training.py
@@ -49,6 +49,7 @@
 from deepmd.dpmodel.train import (
     DEFAULT_TASK_KEY,
     AbstractTrainer,
+    MetricAccumulator,
     RankContext,
     ShardingPolicy,
     TrainerConfig,
@@ -61,6 +62,9 @@
     change_model_out_bias_by_task,
     resolve_step_schedule,
 )
+from deepmd.pt_expt.train.metrics import (
+    all_reduce_metric_accumulator,
+)
 from deepmd.dpmodel.utils.batch import (
     normalize_batch,
     split_batch,
@@ -3278,6 +3282,12 @@ def on_train_begin(self, tasks: TrainingTaskCollection) -> None:
         """Switch the wrapper to training mode."""
         self.wrapper.train()
 
+    def synchronize_metric_accumulator(self) -> None:
+        """All-reduce averaged training metrics before the chief logs them."""
+        if self.metric_accumulator is None or self.world_size <= 1:
+            return
+        all_reduce_metric_accumulator(self.metric_accumulator)
+
     def collect_display_results(
         self,
         tasks: TrainingTaskCollection,
@@ -3382,12 +3392,17 @@ def evaluate_validation(
         Validation is therefore skipped from stage two on; the metrics remain
         available through the independent full validation flow, which every
         rank enters together.
+
+        Display validation uses the shared :class:`MetricAccumulator` with
+        atom weights so optional NaN metrics and empty batches share the same
+        rules as training interval averages.
         """
         if task.validation_data is None or self.sharding.shards_parameters:
             return None
 
-        valid_results: dict[str, float] = {}
-        sum_natoms = 0
+        loss = self.losses[task.key]
+        accumulator = MetricAccumulator({task.key: loss.training_metric_names})
+        saw_batch = False
         for _ii in range(task.valid_numb_batch):
             val_input, val_label = self.get_data(is_train=False, task_key=task.key)
             if not val_input:
@@ -3398,21 +3413,22 @@ def evaluate_validation(
                 label=val_label,
                 task_key=task.key,
             )
-            # The metrics are per-atom quantities, so each batch weighs by the
-            # real atoms it holds summed over its frames. Phantom atoms
-            # (atype < 0), which pad a mixed-nloc batch, contribute to none.
+            # Per-atom metrics weigh by real atoms across frames. Phantom
+            # atoms (atype < 0) pad mixed-nloc batches and contribute nothing.
             natoms = int((val_input["atype"] >= 0).sum())
-            sum_natoms += natoms
-            for key, value in vmore.items():
-                if "l2_" not in key:
-                    valid_results[key] = (
-                        valid_results.get(key, 0.0) + self._to_float(value) * natoms
-                    )
-        if sum_natoms > 0:
-            valid_results = {
-                key: value / sum_natoms for key, value in valid_results.items()
-            }
-        return valid_results
+            accumulator.add(
+                task.key,
+                {
+                    key: value.detach() if torch.is_tensor(value) else value
+                    for key, value in vmore.items()
+                    if "l2_" not in key
+                },
+                weight=float(natoms),
+            )
+            saw_batch = True
+        if not saw_batch:
+            return {}
+        return accumulator.average(task.key)
 
     def learning_rate(self, step: int) -> float:
         """Return the configured learning rate for a zero-based step."""
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 982c5d2a26..f4e748c167 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -6023,8 +6023,12 @@ def training_args(
     doc_disp_avg = (
         "Display the arithmetic mean of per-step training metrics within each "
         "display interval, separately for each task. Tasks with no training "
-        "steps in an interval display NaN. Validation metrics are evaluated "
-        "at the display step, not averaged over training steps."
+        "steps in an interval display NaN. A metric that is NaN on a given "
+        "step (for example an optional label that was absent) is excluded "
+        "from that metric's interval mean. Validation metrics are evaluated "
+        "at the display step with atom weighting, not averaged over training "
+        "steps. Distributed runs reduce local sums and weights only at "
+        "display boundaries."
     )
     doc_profiling = (
         "Enable performance profiling. TensorFlow and PyTorch can export a Chrome "
diff --git a/source/tests/common/dpmodel/test_train_metrics.py b/source/tests/common/dpmodel/test_train_metrics.py
index 8acabcae7a..2ba32172ee 100644
--- a/source/tests/common/dpmodel/test_train_metrics.py
+++ b/source/tests/common/dpmodel/test_train_metrics.py
@@ -7,6 +7,7 @@
 import pytest
 
 from deepmd.dpmodel.train.metrics import (
+    MetricAccumulator,
     TrainingMetricAccumulator,
 )
 
@@ -45,14 +46,53 @@ def test_accumulation_does_not_alias_or_modify_backend_buffers() -> None:
 
 
 @pytest.mark.parametrize("metrics", [{}, {"rmse": float("nan")}])
-def test_missing_and_nan_metrics_are_not_silently_zeroed(metrics: dict) -> None:
-    accumulator = TrainingMetricAccumulator({"task": ("rmse",)})
+def test_missing_and_nan_metrics_are_excluded_from_the_average(metrics: dict) -> None:
+    """Optional labels must not poison an otherwise finite interval average."""
+    accumulator = MetricAccumulator({"task": ("rmse",)})
     accumulator.add("task", {"rmse": 1.0})
     accumulator.add("task", metrics)
+    assert accumulator.average("task") == {"rmse": 1.0}
+    assert accumulator.count("task") == 2
+
+
+def test_metric_with_only_nan_observations_reports_nan() -> None:
+    accumulator = MetricAccumulator({"task": ("rmse", "mae")})
+    accumulator.add("task", {"rmse": float("nan"), "mae": 2.0})
+    accumulator.add("task", {"rmse": float("nan"), "mae": 4.0})
     assert math.isnan(accumulator.average("task")["rmse"])
+    assert accumulator.average("task")["mae"] == 3.0
+
+
+def test_atom_weights_compute_validation_style_averages() -> None:
+    accumulator = MetricAccumulator({"task": ("rmse",)})
+    accumulator.add("task", {"rmse": 1.0}, weight=2.0)
+    accumulator.add("task", {"rmse": 4.0}, weight=1.0)
+    assert accumulator.average("task") == {"rmse": 2.0}
     assert accumulator.count("task") == 2
 
 
+def test_non_positive_weight_skips_metric_totals() -> None:
+    accumulator = MetricAccumulator({"task": ("rmse",)})
+    accumulator.add("task", {"rmse": 1.0}, weight=0.0)
+    assert accumulator.count("task") == 1
+    assert math.isnan(accumulator.average("task")["rmse"])
+
+
+def test_flat_state_round_trip_supports_distributed_reduce() -> None:
+    left = MetricAccumulator({"a": ("rmse",), "b": ("mae",)})
+    right = MetricAccumulator({"a": ("rmse",), "b": ("mae",)})
+    left.add("a", {"rmse": 1.0})
+    left.add("b", {"mae": 2.0})
+    right.add("a", {"rmse": 3.0})
+    # b unsampled on the right rank
+    merged = [x + y for x, y in zip(left.flat_state(), right.flat_state())]
+    left.load_flat_state(merged)
+    assert left.average("a") == {"rmse": 2.0}
+    assert left.average("b") == {"mae": 2.0}
+    assert left.count("a") == 2
+    assert left.count("b") == 1
+
+
 def test_unknown_metric_does_not_partially_update_the_window() -> None:
     accumulator = TrainingMetricAccumulator({"task": ("rmse",)})
     with pytest.raises(ValueError, match=r"Undeclared display metrics.*mae"):
diff --git a/source/tests/pt_expt/train/test_metric_reduce.py b/source/tests/pt_expt/train/test_metric_reduce.py
new file mode 100644
index 0000000000..4f51bf5c68
--- /dev/null
+++ b/source/tests/pt_expt/train/test_metric_reduce.py
@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Tests for Torch metric-window reduction helpers."""
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from deepmd.dpmodel.train.metrics import (  # noqa: E402
+    MetricAccumulator,
+)
+from deepmd.pt_expt.train.metrics import (  # noqa: E402
+    all_reduce_metric_accumulator,
+)
+
+
+def test_all_reduce_is_noop_without_process_group() -> None:
+    accumulator = MetricAccumulator({"task": ("rmse",)})
+    accumulator.add("task", {"rmse": 1.5})
+    all_reduce_metric_accumulator(accumulator)
+    assert accumulator.average("task") == {"rmse": 1.5}

From a7d53407cc9a4b2ee36eae6f08e5a456f1aa6a92 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 9 Oct 2026 15:53:11 +0000
Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/dpmodel/train/metrics.py                  |  4 +++-
 deepmd/dpmodel/train/trainer.py                  |  1 -
 deepmd/pt/train/training.py                      | 10 ++++------
 deepmd/pt_expt/train/metrics.py                  |  4 +++-
 deepmd/pt_expt/train/training.py                 |  6 +++---
 source/tests/pt_expt/train/test_metric_reduce.py |  4 ++--
 6 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/deepmd/dpmodel/train/metrics.py b/deepmd/dpmodel/train/metrics.py
index ab32b1f574..0be3362c51 100644
--- a/deepmd/dpmodel/train/metrics.py
+++ b/deepmd/dpmodel/train/metrics.py
@@ -14,7 +14,9 @@
 weights on device and reduce them only at a display boundary.
 """
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
 from collections.abc import (
     Mapping,
diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py
index 1e22103ca9..6901077b8b 100644
--- a/deepmd/dpmodel/train/trainer.py
+++ b/deepmd/dpmodel/train/trainer.py
@@ -57,7 +57,6 @@
 if TYPE_CHECKING:
     from .metrics import (
         MetricAccumulator,
-        TrainingMetricAccumulator,
     )
 
 DEFAULT_TASK_KEY = "Default"
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index b0200515b3..285a113e0d 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -36,9 +36,6 @@
     change_model_out_bias,
     resolve_step_schedule,
 )
-from deepmd.pt_expt.train.metrics import (
-    all_reduce_metric_accumulator,
-)
 from deepmd.dpmodel.utils import (
     compute_total_numb_batch,
 )
@@ -123,6 +120,9 @@
     NonFiniteGradGuard,
     clip_grad_norm_,
 )
+from deepmd.pt_expt.train.metrics import (
+    all_reduce_metric_accumulator,
+)
 from deepmd.pt_expt.train.utils import (
     count_parameters,
     infer_env_defaults,
@@ -1641,9 +1641,7 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
                             _task_key,
                             {
                                 name: (
-                                    value.detach()
-                                    if torch.is_tensor(value)
-                                    else value
+                                    value.detach() if torch.is_tensor(value) else value
                                 )
                                 for name, value in more_loss.items()
                                 if "l2_" not in name
diff --git a/deepmd/pt_expt/train/metrics.py b/deepmd/pt_expt/train/metrics.py
index c78b7523b8..5a45979ec8 100644
--- a/deepmd/pt_expt/train/metrics.py
+++ b/deepmd/pt_expt/train/metrics.py
@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Torch helpers for distributed metric-window reduction."""
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
 from typing import (
     TYPE_CHECKING,
diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py
index d8fbfe5c85..fb8da802b5 100644
--- a/deepmd/pt_expt/train/training.py
+++ b/deepmd/pt_expt/train/training.py
@@ -62,9 +62,6 @@
     change_model_out_bias_by_task,
     resolve_step_schedule,
 )
-from deepmd.pt_expt.train.metrics import (
-    all_reduce_metric_accumulator,
-)
 from deepmd.dpmodel.utils.batch import (
     normalize_batch,
     split_batch,
@@ -124,6 +121,9 @@
     NonFiniteGradGuard,
     clip_grad_norm_,
 )
+from deepmd.pt_expt.train.metrics import (
+    all_reduce_metric_accumulator,
+)
 from deepmd.pt_expt.train.utils import (
     MatmulPrecisionPolicy,
     count_parameters,
diff --git a/source/tests/pt_expt/train/test_metric_reduce.py b/source/tests/pt_expt/train/test_metric_reduce.py
index 4f51bf5c68..a16bd3fb29 100644
--- a/source/tests/pt_expt/train/test_metric_reduce.py
+++ b/source/tests/pt_expt/train/test_metric_reduce.py
@@ -5,10 +5,10 @@
 
 torch = pytest.importorskip("torch")
 
-from deepmd.dpmodel.train.metrics import (  # noqa: E402
+from deepmd.dpmodel.train.metrics import (
     MetricAccumulator,
 )
-from deepmd.pt_expt.train.metrics import (  # noqa: E402
+from deepmd.pt_expt.train.metrics import (
     all_reduce_metric_accumulator,
 )
 

From 5235a2166878cc7b436fd7cd712f8e6ecf842c7a Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Fri, 9 Oct 2026 23:59:46 +0800
Subject: [PATCH 3/6] fix(train): avoid D2H sync in MetricAccumulator.add

NaN filtering on add only inspects host scalars (None/float/int/0-d
NumPy). Present device metrics stay on-device until average() or a
display-boundary all-reduce; optional-label gaps already arrive as host
nan from display_if_exist.

Refs #5819
---
 deepmd/dpmodel/train/metrics.py               | 44 +++++++++++++------
 .../common/dpmodel/test_train_metrics.py      | 38 ++++++++++++++++
 2 files changed, 69 insertions(+), 13 deletions(-)

diff --git a/deepmd/dpmodel/train/metrics.py b/deepmd/dpmodel/train/metrics.py
index 0be3362c51..4140a64884 100644
--- a/deepmd/dpmodel/train/metrics.py
+++ b/deepmd/dpmodel/train/metrics.py
@@ -8,10 +8,11 @@
 * Validation batches use an atom weight (real atoms in the batch), so the
   reported value is the atom-weighted mean of per-atom metrics.
 
-Optional labels emit NaN through ``display_if_exist``. Those observations are
-excluded from that metric's sum and weight so a single unlabeled batch cannot
-poison an otherwise finite interval. Distributed runs keep local sums and
-weights on device and reduce them only at a display boundary.
+Optional labels emit host NaN through ``display_if_exist``. Those observations
+are excluded from that metric's sum and weight so a single unlabeled batch
+cannot poison an otherwise finite interval. Present device metrics are never
+host-synced on ``add``; distributed runs keep local sums and weights on device
+and reduce them only at a display boundary.
 """
 
 from __future__ import (
@@ -30,14 +31,31 @@
 )
 
 
-def _is_nan(value: Any) -> bool:
-    """Return whether a scalar host or array value is NaN."""
+def _is_host_nan(value: Any) -> bool:
+    """Return whether a host scalar is NaN without touching device memory.
+
+    Optional-label gaps arrive as host ``float("nan")`` from
+    ``display_if_exist``. Present CUDA / accelerator metrics must stay on
+    device through :meth:`MetricAccumulator.add`; coercing them with
+    ``float()`` would D2H-sync on every training step. Device values are
+    therefore never inspected here. Host conversion happens only in
+    :meth:`average` and in display-boundary all-reduce packing.
+    """
     if value is None:
         return True
-    try:
-        return bool(isnan(float(value)))
-    except (TypeError, ValueError):
+    if isinstance(value, bool):
         return False
+    if isinstance(value, (int, float)):
+        return isnan(float(value))
+    # 0-d NumPy arrays and NumPy scalars are already on the host.
+    if type(value).__module__.startswith("numpy"):
+        ndim = getattr(value, "ndim", 0)
+        if ndim == 0:
+            try:
+                return bool(isnan(float(value)))
+            except (TypeError, ValueError):
+                return False
+    return False
 
 
 class MetricAccumulator:
@@ -46,7 +64,7 @@ class MetricAccumulator:
     Metric names are declared before training so an unsampled task keeps the
     same columns as a sampled task. Each call to :meth:`add` is one
     observation: training uses ``weight=1``, validation passes the number of
-    real atoms. A NaN metric is skipped for that observation only.
+    real atoms. A host-NaN metric is skipped for that observation only.
 
     Scalar arrays stay on their original device until :meth:`average`
     converts them to Python floats for display. Out-of-place addition
@@ -89,7 +107,7 @@ def add(
         """Record one observation of detached scalar metrics.
 
         Backends detach metrics from their differentiation graph before
-        calling this method. Missing or NaN metrics are omitted from that
+        calling this method. Missing or host-NaN metrics are omitted from that
         metric's weighted average rather than treated as zero. A non-positive
         weight records the observation for task sampling but contributes to
         no metric totals.
@@ -104,7 +122,7 @@ def add(
         weights = self._weights[task_key]
         for name in names:
             value = metrics.get(name, float("nan"))
-            if weight <= 0.0 or _is_nan(value):
+            if weight <= 0.0 or _is_host_nan(value):
                 continue
             # Out-of-place update: backends may overwrite the tensor buffer
             # that produced ``value`` on the next step.
@@ -127,7 +145,7 @@ def average(self, task_key: str) -> dict[str, float]:
         result: dict[str, float] = {}
         for name in self._names[task_key]:
             weight = float(weights[name])
-            if weight == 0.0 or _is_nan(weight):
+            if weight == 0.0 or isnan(weight):
                 result[name] = float("nan")
             else:
                 result[name] = float(sums[name] / weights[name])
diff --git a/source/tests/common/dpmodel/test_train_metrics.py b/source/tests/common/dpmodel/test_train_metrics.py
index 2ba32172ee..f1b5e2862e 100644
--- a/source/tests/common/dpmodel/test_train_metrics.py
+++ b/source/tests/common/dpmodel/test_train_metrics.py
@@ -63,6 +63,44 @@ def test_metric_with_only_nan_observations_reports_nan() -> None:
     assert accumulator.average("task")["mae"] == 3.0
 
 
+def test_add_does_not_float_coerce_device_scalars() -> None:
+    """Present accelerator metrics must not D2H-sync on every training step."""
+
+    class DeviceScalar:
+        def __init__(self, value: float) -> None:
+            self.value = float(value)
+            self.device = "cuda:0"
+            self.ndim = 0
+            self.float_calls = 0
+
+        def __float__(self) -> float:
+            self.float_calls += 1
+            return self.value
+
+        def __mul__(self, other: object):
+            return DeviceScalar(self.value * float(other))
+
+        __rmul__ = __mul__
+
+        def __add__(self, other: object):
+            if isinstance(other, DeviceScalar):
+                return DeviceScalar(self.value + other.value)
+            return DeviceScalar(self.value + float(other))
+
+        def __radd__(self, other: object):
+            return self.__add__(other)
+
+        def __truediv__(self, other: object):
+            return DeviceScalar(self.value / float(other))
+
+    accumulator = MetricAccumulator({"task": ("rmse",)})
+    device_value = DeviceScalar(2.0)
+    accumulator.add("task", {"rmse": device_value})
+    assert device_value.float_calls == 0
+    # Host conversion is reserved for the display boundary.
+    assert accumulator.average("task") == {"rmse": 2.0}
+
+
 def test_atom_weights_compute_validation_style_averages() -> None:
     accumulator = MetricAccumulator({"task": ("rmse",)})
     accumulator.add("task", {"rmse": 1.0}, weight=2.0)

From 9ade1bc9dc326269034d53f9b79551c1884fcc88 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Sat, 10 Oct 2026 00:05:51 +0800
Subject: [PATCH 4/6] fix(train): pick metric all-reduce device from process
 group

Host-only windows left _state_device on CPU while other ranks used CUDA,
which breaks NCCL collectives. Choose the device from the group backend
and this rank's current CUDA device instead.
---
 deepmd/pt_expt/train/metrics.py               | 21 ++++--
 .../tests/pt_expt/train/test_metric_reduce.py | 69 +++++++++++++++++++
 2 files changed, 84 insertions(+), 6 deletions(-)

diff --git a/deepmd/pt_expt/train/metrics.py b/deepmd/pt_expt/train/metrics.py
index 5a45979ec8..d204153a1a 100644
--- a/deepmd/pt_expt/train/metrics.py
+++ b/deepmd/pt_expt/train/metrics.py
@@ -41,7 +41,10 @@ def all_reduce_metric_accumulator(
     if not packed:
         return
 
-    device = _state_device(packed)
+    # Packing always host-converts via float(); the collective device must
+    # come from the process group, not from whether this rank's window
+    # promoted any sums to tensors.
+    device = _collective_device(group)
     tensor = torch.tensor(
         [float(value) for value in packed],
         dtype=torch.float64,
@@ -51,9 +54,15 @@ def all_reduce_metric_accumulator(
     accumulator.load_flat_state(tensor.detach().cpu().tolist())
 
 
-def _state_device(packed: list[Any]) -> torch.device:
-    """Prefer the device of the first tensor entry; otherwise use CPU."""
-    for value in packed:
-        if torch.is_tensor(value):
-            return value.device
+def _collective_device(group: Any | None = None) -> torch.device:
+    """Return one all-reduce device that every rank in ``group`` can use.
+
+    Device choice follows the process-group backend and this rank's current
+    CUDA device. It must not depend on local accumulator contents: a rank
+    whose window never promoted sums to tensors would otherwise pick CPU
+    while another rank picked CUDA, which breaks NCCL collectives.
+    """
+    backend = str(dist.get_backend(group)).lower()
+    if "nccl" in backend and torch.cuda.is_available():
+        return torch.device("cuda", torch.cuda.current_device())
     return torch.device("cpu")
diff --git a/source/tests/pt_expt/train/test_metric_reduce.py b/source/tests/pt_expt/train/test_metric_reduce.py
index a16bd3fb29..5470e4b239 100644
--- a/source/tests/pt_expt/train/test_metric_reduce.py
+++ b/source/tests/pt_expt/train/test_metric_reduce.py
@@ -1,6 +1,14 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Tests for Torch metric-window reduction helpers."""
 
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    Any,
+)
+
 import pytest
 
 torch = pytest.importorskip("torch")
@@ -8,7 +16,11 @@
 from deepmd.dpmodel.train.metrics import (
     MetricAccumulator,
 )
+from deepmd.pt_expt.train import (
+    metrics as metrics_mod,
+)
 from deepmd.pt_expt.train.metrics import (
+    _collective_device,
     all_reduce_metric_accumulator,
 )
 
@@ -18,3 +30,60 @@ def test_all_reduce_is_noop_without_process_group() -> None:
     accumulator.add("task", {"rmse": 1.5})
     all_reduce_metric_accumulator(accumulator)
     assert accumulator.average("task") == {"rmse": 1.5}
+
+
+def test_collective_device_follows_nccl_not_host_window(monkeypatch: pytest.MonkeyPatch) -> None:
+    """NCCL ranks must agree on CUDA even when the local window is host-only."""
+    monkeypatch.setattr(metrics_mod.dist, "get_backend", lambda group=None: "nccl")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    assert _collective_device() == torch.device("cuda", 0)
+
+    monkeypatch.setattr(
+        metrics_mod.dist,
+        "get_backend",
+        lambda group=None: "cuda:nccl,cpu:gloo",
+    )
+    assert _collective_device() == torch.device("cuda", 0)
+
+
+def test_collective_device_uses_cpu_for_gloo(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(metrics_mod.dist, "get_backend", lambda group=None: "gloo")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    assert _collective_device() == torch.device("cpu")
+
+
+def test_all_reduce_uses_group_device_for_host_only_window(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Host-only flat_state must still all-reduce on the group device."""
+    captured: dict[str, Any] = {}
+
+    def fake_all_reduce(tensor: torch.Tensor, op=None, group=None) -> None:
+        captured["device"] = tensor.device
+        tensor.fill_(0.0)
+
+    monkeypatch.setattr(metrics_mod.dist, "is_available", lambda: True)
+    monkeypatch.setattr(metrics_mod.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(metrics_mod.dist, "get_world_size", lambda group=None: 2)
+    monkeypatch.setattr(metrics_mod.dist, "get_backend", lambda group=None: "nccl")
+    monkeypatch.setattr(metrics_mod.dist, "all_reduce", fake_all_reduce)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    # Keep tensor construction on CPU in this CPU-only CI image while still
+    # asserting that the requested device is the NCCL CUDA device.
+    real_tensor = torch.tensor
+
+    def fake_tensor(data, dtype=None, device=None, **kwargs):
+        captured["requested_device"] = torch.device(device) if device is not None else torch.device("cpu")
+        return real_tensor(data, dtype=dtype, device="cpu", **kwargs)
+
+    monkeypatch.setattr(torch, "tensor", fake_tensor)
+
+    accumulator = MetricAccumulator({"task": ("rmse",)})
+    # Optional-label NaNs leave sums as host floats (never promoted to CUDA).
+    accumulator.add("task", {"rmse": float("nan")})
+    assert all(not torch.is_tensor(v) for v in accumulator.flat_state())
+
+    all_reduce_metric_accumulator(accumulator)
+    assert captured["requested_device"] == torch.device("cuda", 0)

From b51e6ed6a9470e7ca47ebbe665af4b3b4e8fa5d1 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 9 Oct 2026 16:06:54 +0000
Subject: [PATCH 5/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/pt_expt/train/test_metric_reduce.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/source/tests/pt_expt/train/test_metric_reduce.py b/source/tests/pt_expt/train/test_metric_reduce.py
index 5470e4b239..5c1d07921c 100644
--- a/source/tests/pt_expt/train/test_metric_reduce.py
+++ b/source/tests/pt_expt/train/test_metric_reduce.py
@@ -16,9 +16,7 @@
 from deepmd.dpmodel.train.metrics import (
     MetricAccumulator,
 )
-from deepmd.pt_expt.train import (
-    metrics as metrics_mod,
-)
+from deepmd.pt_expt.train import metrics as metrics_mod
 from deepmd.pt_expt.train.metrics import (
     _collective_device,
     all_reduce_metric_accumulator,
@@ -32,7 +30,9 @@ def test_all_reduce_is_noop_without_process_group() -> None:
     assert accumulator.average("task") == {"rmse": 1.5}
 
 
-def test_collective_device_follows_nccl_not_host_window(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_collective_device_follows_nccl_not_host_window(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     """NCCL ranks must agree on CUDA even when the local window is host-only."""
     monkeypatch.setattr(metrics_mod.dist, "get_backend", lambda group=None: "nccl")
     monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
@@ -75,7 +75,9 @@ def fake_all_reduce(tensor: torch.Tensor, op=None, group=None) -> None:
     real_tensor = torch.tensor
 
     def fake_tensor(data, dtype=None, device=None, **kwargs):
-        captured["requested_device"] = torch.device(device) if device is not None else torch.device("cpu")
+        captured["requested_device"] = (
+            torch.device(device) if device is not None else torch.device("cpu")
+        )
         return real_tensor(data, dtype=dtype, device="cpu", **kwargs)
 
     monkeypatch.setattr(torch, "tensor", fake_tensor)

From 0ff57049611d8b1759c1f067ec4c973626d47045 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <198997215+OutisLi-Bot@users.noreply.github.com>
Date: Sat, 10 Oct 2026 10:59:12 +0800
Subject: [PATCH 6/6] style(train): silence TC003/B905 on MetricAccumulator

Move Mapping/Sequence under TYPE_CHECKING and mark flat_state zip strict.
---
 deepmd/dpmodel/train/metrics.py                   | 11 +++++++----
 source/tests/common/dpmodel/test_train_metrics.py |  2 +-
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/deepmd/dpmodel/train/metrics.py b/deepmd/dpmodel/train/metrics.py
index 4140a64884..2996e053d6 100644
--- a/deepmd/dpmodel/train/metrics.py
+++ b/deepmd/dpmodel/train/metrics.py
@@ -19,17 +19,20 @@
     annotations,
 )
 
-from collections.abc import (
-    Mapping,
-    Sequence,
-)
 from math import (
     isnan,
 )
 from typing import (
+    TYPE_CHECKING,
     Any,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import (
+        Mapping,
+        Sequence,
+    )
+
 
 def _is_host_nan(value: Any) -> bool:
     """Return whether a host scalar is NaN without touching device memory.
diff --git a/source/tests/common/dpmodel/test_train_metrics.py b/source/tests/common/dpmodel/test_train_metrics.py
index f1b5e2862e..ab6a34a450 100644
--- a/source/tests/common/dpmodel/test_train_metrics.py
+++ b/source/tests/common/dpmodel/test_train_metrics.py
@@ -123,7 +123,7 @@ def test_flat_state_round_trip_supports_distributed_reduce() -> None:
     left.add("b", {"mae": 2.0})
     right.add("a", {"rmse": 3.0})
     # b unsampled on the right rank
-    merged = [x + y for x, y in zip(left.flat_state(), right.flat_state())]
+    merged = [x + y for x, y in zip(left.flat_state(), right.flat_state(), strict=True)]
     left.load_flat_state(merged)
     assert left.average("a") == {"rmse": 2.0}
     assert left.average("b") == {"mae": 2.0}