deepmodeling / deepmodeling/deepmd-kit
feat(pt_expt): honor training.save_dir through a checkpoint manager
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Problem
pt_expt accepts `training.save_dir`, but checkpoint paths are constructed directly from `save_ckpt`, so the configured directory is ignored. Periodic checkpoints, the latest alias, retention cleanup, and restart discovery can consequently operate in different locations.
## Root cause
Checkpoint path resolution, serialization, latest-link management, retention, and restart discovery are spread across trainer methods. There is no single owner of checkpoint naming or filesystem state.
## Proposed architecture
Introduce a pt_expt-owned `TorchCheckpointManager`, with backend-neutral path and retention policy factored under `deepmd/dpmodel/train` where useful.
The manager should own:
- canonical resolution of `save_dir`, `save_ckpt`, and step filenames;
- directory creation;
- atomic checkpoint publication using a temporary file followed by `os.replace`;
- latest-checkpoint alias or metadata updates;
- retention for regular and EMA checkpoint namespaces;
- explicit and latest-checkpoint restart resolution;
- rank ownership of filesystem side effects.
Preserve existing PT path semantics when `save_dir` is set, and existing prefix-relative behavior when it is absent. Trainer code should provide checkpoint payloads; it should not assemble filenames or scan directories itself.
## Correctness and performance requirements
- Only the chief rank may create, replace, link, or delete checkpoint files.
- Failed or interrupted serialization must not replace the previous latest checkpoint.
- Regular, EMA, and validation-best checkpoints must have independent namespaces and retention policies.
- Retention must never delete the latest alias target or unrelated files.
- Avoid unconditional distributed barriers and per-step directory scans; filesystem work should occur only at checkpoint boundaries.
## Acceptance criteria
- With `save_dir="./models"`, all periodic checkpoints and latest metadata are created under `./models`.
- Omitting `save_dir` preserves current behavior.
- Nested and absolute `save_ckpt` inputs have documented, tested resolution.
- `max_ckpt_keep` removes only the intended old step checkpoints.
- Restart from both an explicit checkpoint and the latest checkpoint works with `save_dir`.
- An interrupted save leaves the previous checkpoint and latest reference valid.
- DDP tests confirm that only rank zero performs filesystem mutations.
- Trainer checkpoint methods delegate path and retention behavior to the manager.
Refs #5755.
Contributor guide
Assessment
This issue has not been assessed yet.