deepmodeling / deepmodeling/deepmd-kit

feat(train): support interval metric aggregation for disp_avg

Open
#5,819 1 comment 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

## 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 #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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.