deepmodeling / deepmodeling/deepmd-kit

feat(pt_expt): honor training.save_dir through a checkpoint manager

Open
#5,817 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

## 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

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.