deepmodeling / deepmodeling/unimol_tools
[Code scan] Aggregate pretraining DDP validation metrics across ranks
- Dominant language
- Python
- Stars
- 34
- Forks
- 9
- PR merge metrics
- No merged PRs in 30d
Description
This issue is a result of a Codex global repository scan.
## Summary
In pretraining DDP evaluation, each rank evaluates its validation shard, but rank 0 chooses best-checkpoint and early-stopping decisions from its local metrics. The metric reducer multiplies bsz by world size, but loss sums and sample counts are not all-reduced or gathered before checkpoint decisions.
## Code references
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/trainer.py#L296-L321
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/trainer.py#L400-L417
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/trainer.py#L442-L458
## Impact
Best-checkpoint selection and early stopping can depend on only rank 0 validation samples, not the full validation set. This can select the wrong checkpoint or stop too early/late in distributed pretraining.
## Suggested fix
All-reduce metric numerators and counts, or gather logging outputs from all ranks, before reduce_metrics and before checkpoint or early-stop decisions.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with unimol_tools/pretrain/trainer.py at lines 296-321, 400-417, and 442-458. Trace how validation loss sums, sample counts, reduce_metrics, checkpoint selection, and early stopping use rank-local values. Done means validation metrics are aggregated across ranks before reduction and before checkpoint or early-stop decisions.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100