feat(train): shared MetricAccumulator for disp_avg intervals

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",

deepmd/dpmodel/train/metrics.py

@@ -1,23 +1,77 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""Backend-independent accumulation of training display metrics."""
+"""Backend-independent accumulation of training display metrics.
 
-from collections.abc import (
-    Mapping,
-    Sequence,
+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 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 (
+    annotations,
+)
+
+from math import (
+    isnan,
 )
 from typing import (
+    TYPE_CHECKING,
     Any,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import (
+        Mapping,
+        Sequence,
+    )
+
 
-class TrainingMetricAccumulator:
-    """Average per-step metrics independently for each training task.
+def _is_host_nan(value: Any) -> bool:
+    """Return whether a host scalar is NaN without touching device memory.
 
-    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.
+    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
+    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:
+    """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 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
+    preserves recorded values when a backend reuses an output buffer.
 
     Parameters
     ----------
@@ -26,44 +80,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._weights: 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._observations = dict.fromkeys(self._names, 0)
 
-    def add(self, task_key: str, metrics: Mapping[str, Any]) -> None:
-        """Record one optimizer step's detached scalar metrics.
+    @property
+    def task_keys(self) -> tuple[str, ...]:
+        """Return the declared task keys in insertion order."""
+        return tuple(self._names)
+
+    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 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.
         """
-        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_host_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 isnan(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

deepmd/dpmodel/train/trainer.py

@@ -56,7 +56,7 @@
 
 if TYPE_CHECKING:
     from .metrics import (
-        TrainingMetricAccumulator,
+        MetricAccumulator,
     )
 
 DEFAULT_TASK_KEY = "Default"
@@ -537,7 +537,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 +572,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 +669,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

deepmd/pt/train/training.py

@@ -28,6 +28,7 @@
 from deepmd.dpmodel.train import (
     DEFAULT_TASK_KEY,
     CheckpointStore,
+    MetricAccumulator,
     ShardingPolicy,
     TrainingMetricAccumulator,
     TrainingTimer,
@@ -119,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,
@@ -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,24 @@ 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)