fix(disp): div 0 error when using disp_avg under multitask

deepmd/dpmodel/loss/dos.py

@@ -231,6 +231,21 @@ def call(
         more_loss["rmse"] = xp.sqrt(loss)
         return loss, more_loss
 
+    @property
+    def training_metric_names(self) -> tuple[str, ...]:
+        """Return configured global and atomic DOS/CDF metrics."""
+        names = tuple(
+            name
+            for name, enabled in (
+                ("rmse_global_dos", self.has_dos),
+                ("rmse_global_cdf", self.has_cdf),
+                ("rmse_local_dos", self.has_ados),
+                ("rmse_local_cdf", self.has_acdf),
+            )
+            if enabled
+        )
+        return ("rmse", *names)
+
     @property
     def label_requirement(self) -> list[DataRequirementItem]:
         """Return data label requirements needed for this loss calculation."""

deepmd/dpmodel/loss/ener.py

@@ -876,6 +876,20 @@ def call(
         self.l2_more = more_loss
         return loss, more_loss
 
+    @property
+    def training_metric_names(self) -> tuple[str, ...]:
+        """Return configured energy, force, virial and Hessian metrics."""
+        prefix = "rmse" if self.loss_func == "mse" else "mae"
+        names = ("rmse",)
+        names += tuple(
+            f"{prefix}_{term}"
+            for term in ("e", "f", "v", "ae", "pf")
+            if getattr(self, f"has_{term}")
+        )
+        return names + tuple(
+            f"rmse_{term}" for term in ("gf", "h") if getattr(self, f"has_{term}")
+        )
+
     @property
     def label_requirement(self) -> list[DataRequirementItem]:
         """Return data label requirements needed for this loss calculation."""

deepmd/dpmodel/loss/ener_spin.py

@@ -437,6 +437,17 @@ def reshape_atomic(value: Array, ncomp: int) -> Array:
         more_loss["rmse"] = xp.sqrt(loss)
         return loss, more_loss
 
+    @property
+    def training_metric_names(self) -> tuple[str, ...]:
+        """Return configured energy and real/magnetic force metrics."""
+        prefix = "rmse" if self.loss_func == "mse" else "mae"
+        names = ("rmse",)
+        return names + tuple(
+            f"{prefix}_{term}"
+            for term in ("e", "fr", "fm", "v", "ae")
+            if getattr(self, f"has_{term}")
+        )
+
     @property
     def label_requirement(self) -> list[DataRequirementItem]:
         """Return data label requirements needed for this loss calculation."""

deepmd/dpmodel/loss/loss.py

@@ -46,6 +46,17 @@ def call(
             Additional loss terms/metrics for logging.
         """
 
+    @property
+    def training_metric_names(self) -> tuple[str, ...]:
+        """Display columns for the default training call, excluding l2 terms.
+
+        The names depend on the loss configuration, not on whether a task or
+        label has been sampled. Averaged logging uses them without a forward.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} must define training_metric_names for disp_avg."
+        )
+
     @property
     @abstractmethod
     def label_requirement(self) -> list[DataRequirementItem]:

deepmd/dpmodel/loss/property.py

@@ -179,6 +179,15 @@ def call(
 
         return loss, more_loss
 
+    @property
+    def training_metric_names(self) -> tuple[str, ...]:
+        """Return the selected property metrics."""
+        return tuple(
+            name
+            for name in ("smooth_mae", "mae", "mse", "rmse", "mape")
+            if name in self.metric
+        )
+
     @property
     def label_requirement(self) -> list[DataRequirementItem]:
         """Return data label requirements needed for this loss calculation."""

deepmd/dpmodel/loss/tensor.py

@@ -164,6 +164,19 @@ def call(
         more_loss["rmse"] = xp.sqrt(loss)
         return loss, more_loss
 
+    @property
+    def training_metric_names(self) -> tuple[str, ...]:
+        """Return configured local and global tensor metrics."""
+        names = tuple(
+            f"rmse_{scope}_{self.tensor_name}"
+            for scope, enabled in (
+                ("local", self.has_local_weight),
+                ("global", self.has_global_weight),
+            )
+            if enabled
+        )
+        return ("rmse", *names)
+
     @property
     def label_requirement(self) -> list[DataRequirementItem]:
         """Return data label requirements needed for this loss calculation."""

deepmd/dpmodel/train/__init__.py

@@ -16,6 +16,9 @@
     AbstractTrainEntrypoint,
     TrainEntrypointOptions,
 )
+from .metrics import (
+    TrainingMetricAccumulator,
+)
 from .schedule import (
     StepSchedule,
     resolve_step_schedule,
@@ -53,6 +56,7 @@
     "TrainEntrypointOptions",
     "TrainStepResult",
     "TrainerConfig",
+    "TrainingMetricAccumulator",
     "TrainingTask",
     "TrainingTaskCollection",
     "TrainingTaskConfig",

deepmd/dpmodel/train/metrics.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Backend-independent accumulation of training display metrics."""
+
+from collections.abc import (
+    Mapping,
+    Sequence,
+)
+from typing import (
+    Any,
+)
+
+
+class TrainingMetricAccumulator:
+    """Average per-step metrics independently for each training task.
+
+    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.
+
+    Parameters
+    ----------
+    metric_names : Mapping[str, Sequence[str]]
+        Display metric names for each task, excluding internal loss terms.
+    """
+
+    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._counts = dict.fromkeys(metric_names, 0)
+
+    def add(self, task_key: str, metrics: Mapping[str, Any]) -> None:
+        """Record one optimizer step's 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.
+        """
+        totals = self._totals[task_key]
+        unknown = metrics.keys() - totals.keys()
+        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
+
+    def count(self, task_key: str) -> int:
+        """Return the number of recorded optimizer steps for one task."""
+        return self._counts[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()}
+
+    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

deepmd/dpmodel/train/trainer.py

@@ -34,6 +34,7 @@
     Path,
 )
 from typing import (
+    TYPE_CHECKING,
     Any,
     TextIO,
 )
@@ -53,6 +54,11 @@
     TrainingTimer,
 )
 
+if TYPE_CHECKING:
+    from .metrics import (
+        TrainingMetricAccumulator,
+    )
+
 DEFAULT_TASK_KEY = "Default"
 
 log = logging.getLogger(__name__)
@@ -310,12 +316,16 @@ def _normalize_probabilities(
 
 @dataclass
 class TrainStepResult:
-    """Backend payload returned from one optimizer step."""
+    """Backend payload returned from one optimizer step.
+
+    Averaged display consumes ``train_results`` as detached scalar metrics.
+    These may remain backend arrays until they are converted for display.
+    """
 
     task_key: str
     step: int
     payload: Any = None
-    train_results: LossResults | None = None
+    train_results: Mapping[str, Any] | None = None
 
 
 class LearningCurveWriter:
@@ -361,6 +371,7 @@ def log_results(
         learning_rate: float,
         train_results: DisplayResults,
         valid_results: DisplayResults | None,
+        unsampled_tasks: Sequence[str] = (),
     ) -> None:
         """Log per-task loss results."""
         if self._is_multitask(train_results):
@@ -379,6 +390,7 @@ def log_results(
                         task_name=f"{task_key}_trn",
                         rmse=task_train,
                         learning_rate=learning_rate,
+                        check_total_rmse_nan=task_key not in unsampled_tasks,
                     )
                 )
                 task_valid = valid_task_results.get(task_key)
@@ -399,6 +411,7 @@ def log_results(
                     task_name="trn",
                     rmse=train_results,
                     learning_rate=learning_rate,
+                    check_total_rmse_nan=not unsampled_tasks,
                 )
             )
             if valid_results:
@@ -513,17 +526,22 @@ class AbstractTrainer(ABC):
 
     Backend trainers implement one optimizer step, metric evaluation, learning
     rate lookup, and checkpoint persistence.  This base class handles the
-    common training loop around those hooks.
+    common training loop around those hooks. Backends opt into interval
+    averaging by supplying a ``metric_accumulator`` and detached
+    ``TrainStepResult.train_results``; display then uses those values instead
+    of calling ``evaluate_training``.
     """
 
     def __init__(
         self,
         trainer_config: TrainerConfig,
         *,
         rank_context: RankContext | None = None,
+        metric_accumulator: TrainingMetricAccumulator | None = None,
     ) -> None:
         self.trainer_config = trainer_config
         self.rank_context = rank_context or RankContext()
+        self.metric_accumulator = metric_accumulator
         self.lcurve_writer = LearningCurveWriter()
 
     def run(self, tasks: TrainingTaskCollection) -> None:
@@ -533,6 +551,8 @@ def run(self, tasks: TrainingTaskCollection) -> None:
         num_steps = self.trainer_config.num_steps
         fout: TextIO | None = None
         try:
+            if self.metric_accumulator is not None:
+                self.metric_accumulator.reset()
             self.on_train_begin(tasks)
             fout = self._open_learning_curve()
             timer = TrainingTimer(
@@ -543,6 +563,12 @@ def run(self, tasks: TrainingTaskCollection) -> None:
             for step in range(start_step, num_steps):
                 task = self.select_task(tasks)
                 step_result = self.train_step(task, step)
+                if self.metric_accumulator is not None:
+                    if step_result.train_results is None:
+                        raise RuntimeError(
+                            "Averaged display requires train_step to return metrics."
+                        )
+                    self.metric_accumulator.add(task.key, step_result.train_results)
                 display_step = step + 1
 
                 if self._should_display(display_step):
@@ -559,6 +585,15 @@ def run(self, tasks: TrainingTaskCollection) -> None:
                             learning_rate=current_lr,
                             train_results=train_results,
                             valid_results=valid_results,
+                            unsampled_tasks=(
+                                [
+                                    key
+                                    for key in tasks.keys
+                                    if self.metric_accumulator.count(key) == 0
+                                ]
+                                if self.metric_accumulator is not None
+                                else ()
+                            ),
                         )
                         self._log_interval(timer.record(display_step))
                         if fout is not None:
@@ -575,6 +610,8 @@ def run(self, tasks: TrainingTaskCollection) -> None:
                                 train_results=train_results,
                                 valid_results=valid_results,
                             )
+                    if self.metric_accumulator is not None:
+                        self.metric_accumulator.reset()
 
                 self.run_full_validation(
                     step=step,
@@ -610,26 +647,22 @@ def collect_display_results(
         step_result: TrainStepResult,
     ) -> tuple[DisplayResults, DisplayResults | None]:
         """Collect training and validation results for display."""
-        if not tasks.is_multitask:
-            return (
-                self.evaluate_training(active_task, step, step_result),
-                self.evaluate_validation(active_task, step, step_result),
-            )
-
         train_results: TaskResults = {}
         valid_results: TaskResults = {}
         for task in tasks:
             task_step_result = step_result if task.key == active_task.key else None
-            train_results[task.key] = self.evaluate_training(
-                task,
-                step,
-                task_step_result,
+            train_results[task.key] = (
+                self.metric_accumulator.average(task.key)
+                if self.metric_accumulator is not None
+                else self.evaluate_training(task, step, task_step_result)
             )
             valid_results[task.key] = self.evaluate_validation(
                 task,
                 step,
                 task_step_result,
             )
+        if not tasks.is_multitask:
+            return train_results[active_task.key], valid_results[active_task.key]
         return train_results, valid_results
 
     def on_train_begin(self, tasks: TrainingTaskCollection) -> None:

deepmd/pt/loss/denoise.py

@@ -42,6 +42,19 @@ def __init__(
         self.mask_loss_coord = mask_loss_coord
         self.mask_loss_token = mask_loss_token
 
+    @property
+    def training_metric_names(self) -> tuple[str, ...]:
+        """Return the enabled denoising metrics."""
+        return tuple(
+            name
+            for name, enabled in (
+                ("coord_l1_error", self.has_coord),
+                ("token_error", self.has_token),
+                ("norm_loss", self.has_norm),
+            )
+            if enabled
+        )
+
     def forward(
         self,
         model_pred: dict[str, torch.Tensor],