Skip to content

feat(train): support interval metric aggregation for disp_avg #5819

Description

@OutisLi

Summary

Implement training.disp_avg in the backend-independent training loop and enable it for pt_expt.

pt_expt currently accepts the option but reports only the latest step metrics. The legacy PT trainer has separate interval-aggregation logic, which should be replaced by one shared implementation.

Root cause

The common trainer loop has no first-class metric aggregation policy. Consequently, each backend must either implement averaging independently or silently ignore disp_avg.

Proposed design

Add a common MetricAccumulator abstraction that:

  • accumulates detached sums and counts per task and metric;
  • distinguishes training step averages from atom-weighted validation averages;
  • supports sparse multi-task sampling without dividing by zero;
  • resets only after a display interval is emitted;
  • optionally reduces sums and counts across ranks at display time.

Backend-neutral aggregation should live under deepmd/dpmodel/train. Torch-specific tensor/reduction helpers may live under deepmd/pt_expt/train, with legacy PT importing the same implementation.

Performance and correctness requirements

  • Keep accumulated values on device where possible.
  • Perform host synchronization only at display boundaries.
  • Do not fetch extra training batches merely to display inactive tasks.
  • Do not introduce operations into the compiled model graph.
  • Preserve existing behavior when disp_avg=false.
  • Preserve the unsampled-task safeguards covered by fix(disp): div 0 error when using disp_avg under multitask #5809.

Acceptance criteria

  • disp_avg=true reports the mean over the completed display interval.
  • disp_avg=false reports the existing instantaneous metrics.
  • Single-task and sparsely sampled multi-task cases are covered.
  • No NaN or division by zero occurs when a task has zero samples in an interval.
  • The accumulator resets correctly after regular and final partial intervals.
  • Distributed aggregation uses global sums and counts with collectives only at display time.

Refs #5755.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions