deepmodeling / deepmodeling/deepmd-kit

feat(tf2): support epoch-based and multi-task training horizons

Open
#5,879 0 comments 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

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

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.