deepmodeling / deepmodeling/deepmd-kit
feat(tf2): support epoch-based and multi-task training horizons
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Summary
The common schema accepts `training.numb_epoch` and aliases plus multi-task `training.num_epoch_dict`, but TF2 training still initializes directly from `training.numb_steps`. Epoch-only inputs therefore do not have a complete TF2 execution path.
## Scope
- Resolve single-task epoch counts after dataset and batch-size setup determines steps per epoch.
- Resolve multi-task epoch dictionaries into documented sampling probabilities and a deterministic total step count.
- Define epoch transitions, reshuffling, and restart state.
- Keep distributed ranks aligned on optimizer-step count.
- Preserve explicit fixed-step behavior.
Use a shared backend-neutral horizon resolver where feasible so schema semantics remain consistent.
## Acceptance criteria
- A normalized TF2 input with only an epoch count trains without requiring `numb_steps`.
- Multi-task epoch dictionaries produce documented and tested scheduling behavior.
- Filtering and batch-size resolution occur before horizon calculation.
- Restart resumes the same epoch and intra-epoch position.
- Single-rank and distributed tests cover reshuffling and uneven tails.
- Existing fixed-step TF2 training remains unchanged.
Refs #5757.
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Contributor guide
Assessment
This issue has not been assessed yet.