deepmodeling / deepmodeling/unimol_tools

[Code scan] Aggregate pretraining DDP validation metrics across ranks

Open
#38 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.