deepmodeling / deepmodeling/deepmd-kit
feat(train): support interval metric aggregation for disp_avg
- 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
Assessment
This issue has not been assessed yet.