deepmodeling / deepmodeling/deepmd-kit
feat(pt_expt): support epoch-based training horizons
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Current status
The epoch-horizon and distributed-data parts of the original issue are implemented on `master`:
- `training.numb_epoch` and `num_epochs` are resolved without requiring `numb_steps`;
- epoch cardinality is resolved after batch-size and dataset filtering;
- LMDB batches are deterministically reshuffled by epoch and sharded across ranks with equal optimizer-step counts;
- chief-only statistics initialization and model-stat synchronization are implemented;
- explicit `numb_steps` behavior remains available.
The remaining functional gap is **checkpointable data-source progress**. This also consolidates the only unfinished part of #5823.
## Remaining problem
A training restart restores the model, optimizer, learning-rate schedule, global step, and EMA state, but it does not restore the data source to the same epoch and intra-epoch cursor. LMDB and directory/HDF5 sources therefore restart their iteration state instead of yielding the same next batch as an uninterrupted run.
The global training step is not sufficient to reconstruct this state in general because batch filtering, same-`nloc` grouping, distributed tail policy, multi-task selection, and data-source-specific iteration can change the mapping from optimizer step to batch.
## Root cause
The training data-source interface exposes epoch length and batch iteration, but it does not expose a serializable progress state. Consequently:
- checkpoints contain no epoch/cursor or sampler state;
- restart cannot restore the exact next batch;
- iterators or prefetchers may be constructed before progress is restored;
- multi-task training has no canonical place to persist per-task data progress.
## Proposed design
Add a backend-neutral progress contract under `deepmd/dpmodel/train`, implemented by both LMDB and directory/HDF5 adapters:
- `state_dict()` returns the minimal logical progress state;
- `load_state_dict(state)` restores it before iterator or prefetch construction;
- the state records epoch and intra-epoch cursor, plus only the deterministic sampler/RNG state that cannot be derived from those values;
- multi-task checkpoints store progress independently for each task/data source;
- distributed restart restores one consistent global epoch/cursor and deterministically derives each rank's shard.
Checkpoint payloads should store logical progress, not decoded batches, full sample-index lists, or dataset contents. Older checkpoints without data-source state must remain loadable and retain the current restart behavior.
## Performance requirements
- No per-batch distributed collective or host synchronization.
- No serialization of full permutations, samples, or LMDB contents.
- Restoring progress must not scan or decode the dataset.
- Rank-local shards remain deterministic and keep equal optimizer-step counts.
- Prefetching, if enabled, must not advance the persisted logical cursor past the last completed optimizer step.
## Acceptance criteria
- Existing explicit `numb_steps` and `numb_epoch` horizon behavior remains unchanged.
- Restarting in the middle of an epoch yields the same next batch as an uninterrupted run for LMDB and directory/HDF5 data.
- Restarting at an epoch boundary preserves deterministic reshuffling.
- Distributed ranks restore the same global progress and resume disjoint deterministic shards.
- Multi-task training restores each task's data progress and task-selection state consistently.
- Checkpoints without data-source progress remain backward compatible.
Refs #5755.
Contributor guide
Assessment
This issue has not been assessed yet.