feat(pt): add ckpt_keep_ratio to set max_ckpt_keep automatically

Author: OutisLi

deepmd/pt/train/training.py

@@ -77,6 +77,7 @@
     clip_grad_norm_,
     latest_checkpoint_path,
     resolve_best_checkpoint_dir,
+    resolve_keep_ckpt_count,
     scoped_env_defaults,
 )
 from deepmd.pt.train.validation import (
@@ -216,6 +217,7 @@ def __init__(
             self.save_dir.mkdir(parents=True, exist_ok=True)
         self.save_freq = training_params.get("save_freq", 1000)
         self.max_ckpt_keep = training_params.get("max_ckpt_keep", 5)
+        self.ckpt_keep_ratio = training_params.get("ckpt_keep_ratio")
         self.enable_ema = bool(training_params.get("enable_ema", False))
         self.ema_decay = float(training_params.get("ema_decay", 0.999))
         self.ema_ckpt_keep = int(training_params.get("ema_ckpt_keep", 3))
@@ -729,6 +731,24 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                     rank=self.rank,
                 )
 
+        # === Derive checkpoint retention from ckpt_keep_ratio ===
+        # num_steps is final here (including when derived from num_epoch), so the
+        # ratio can be converted into an absolute keep count once.
+        keep_ckpt_count = resolve_keep_ckpt_count(
+            self.ckpt_keep_ratio, self.num_steps, self.save_freq
+        )
+        if keep_ckpt_count is not None:
+            self.max_ckpt_keep = keep_ckpt_count
+            self.ema_ckpt_keep = keep_ckpt_count
+            log.info(
+                "Resolved checkpoint retention to %d from ckpt_keep_ratio=%s "
+                "(num_steps=%d, save_freq=%d).",
+                keep_ckpt_count,
+                self.ckpt_keep_ratio,
+                self.num_steps,
+                self.save_freq,
+            )
+
         # Learning rate
         self.gradient_max_norm = training_params.get("gradient_max_norm", 0.0)
         self.nonfinite_grad_guard = NonFiniteGradGuard()

deepmd/pt/train/utils.py

@@ -9,6 +9,9 @@
 from contextlib import (
     contextmanager,
 )
+from math import (
+    ceil,
+)
 from pathlib import (
     Path,
 )
@@ -244,3 +247,38 @@ def resolve_best_checkpoint_dir(
     if save_best_dir:
         return Path(save_best_dir)
     return Path(save_ckpt).parent
+
+
+def resolve_keep_ckpt_count(
+    ckpt_keep_ratio: float | None, num_steps: int, save_freq: int
+) -> int | None:
+    """
+    Convert a checkpoint-retention ratio into a sliding-window keep count.
+
+    Periodic checkpoints are saved every ``save_freq`` steps, so a run of
+    ``num_steps`` produces about ``num_steps // save_freq`` of them. Keeping the
+    most recent ``ceil(ratio * total)`` of those is equivalent to retaining the
+    final ``ratio`` fraction of the run by step, without the caller computing
+    the count by hand.
+
+    Parameters
+    ----------
+    ckpt_keep_ratio : float or None
+        The fraction of the training run, by step, whose periodic checkpoints
+        are retained. ``None`` leaves the keep count unchanged.
+    num_steps : int
+        The total number of training steps, already resolved (including when
+        derived from ``numb_epoch``).
+    save_freq : int
+        The checkpoint saving frequency in steps.
+
+    Returns
+    -------
+    int or None
+        The number of most recent checkpoints to keep (at least one), or
+        ``None`` when ``ckpt_keep_ratio`` is not set.
+    """
+    if ckpt_keep_ratio is None:
+        return None
+    total_periodic_ckpts = max(1, num_steps // save_freq)
+    return max(1, ceil(ckpt_keep_ratio * total_periodic_ckpts))

deepmd/utils/argcheck.py

@@ -4994,6 +4994,12 @@ def training_args(
         "The oldest checkpoints will be deleted once the number of checkpoints exceeds max_ckpt_keep. "
         "Defaults to 5."
     )
+    doc_ckpt_keep_ratio = (
+        "An alternative to `max_ckpt_keep` that sets the number of retained "
+        "checkpoints as a fraction in (0, 1) of the run: the most recent "
+        "`ceil(ckpt_keep_ratio * numb_steps / save_freq)` checkpoints are kept. "
+        "When set, it overrides `max_ckpt_keep` and `ema_ckpt_keep`."
+    )
     doc_enable_ema = (
         "Whether to maintain an exponential moving average (EMA) of model "
         "parameters during training and save periodic EMA checkpoints with an "
@@ -5129,6 +5135,15 @@ def training_args(
             "save_ckpt", str, optional=True, default="model.ckpt", doc=doc_save_ckpt
         ),
         Argument("max_ckpt_keep", int, optional=True, default=5, doc=doc_max_ckpt_keep),
+        Argument(
+            "ckpt_keep_ratio",
+            [float, None],
+            optional=True,
+            default=None,
+            doc=doc_only_pt_supported + doc_ckpt_keep_ratio,
+            extra_check=lambda x: x is None or 0.0 < x < 1.0,
+            extra_check_errmsg="must be a fraction in the open interval (0, 1)",
+        ),
         Argument(
             "enable_ema",
             bool,

doc/train/training-advanced.md

@@ -104,6 +104,7 @@ Other keys in the {ref}`training <training>` section are explained below:
 - {ref}`disp_freq <training/disp_freq>` The frequency of printing learning curve. Set in the unit of training steps
 - {ref}`save_freq <training/save_freq>` The frequency of saving checkpoint.
 - {ref}`save_dir <training/save_dir>` The directory where periodic checkpoints are written (PyTorch backend). It is created recursively if missing, while the `model.ckpt.pt` symlinks and the `checkpoint` pointer file stay in the working directory. Defaults to the working directory.
+- {ref}`ckpt_keep_ratio <training/ckpt_keep_ratio>` An alternative to `max_ckpt_keep` (PyTorch backend) that keeps a sliding window of `ceil(ckpt_keep_ratio * numb_steps / save_freq)` most recent checkpoints, i.e. the final `ckpt_keep_ratio` fraction of the run by step. It overrides `max_ckpt_keep` (and `ema_ckpt_keep`) when set, and works the same whether the run length is given by `numb_steps` or `numb_epoch`.
 
 ## Options and environment variables
 

source/tests/pt/test_train_utils.py

@@ -6,6 +6,7 @@
 from deepmd.pt.train.utils import (
     NonFiniteGradGuard,
     clip_grad_norm_,
+    resolve_keep_ckpt_count,
 )
 
 
@@ -111,5 +112,22 @@ def test_resets_after_check(self) -> None:
         guard.raise_if_nonfinite(self._named(1.0))
 
 
+class TestResolveKeepCkptCount(unittest.TestCase):
+    def test_none_ratio_leaves_count_unchanged(self) -> None:
+        self.assertIsNone(resolve_keep_ckpt_count(None, 1000, 10))
+
+    def test_ratio_maps_to_recent_window_count(self) -> None:
+        # 1000 / 10 = 100 periodic checkpoints; 40% keeps the most recent 40.
+        self.assertEqual(resolve_keep_ckpt_count(0.4, 1000, 10), 40)
+
+    def test_ratio_rounds_up(self) -> None:
+        # 4 periodic checkpoints; ceil(0.4 * 4) = ceil(1.6) = 2.
+        self.assertEqual(resolve_keep_ckpt_count(0.4, 4, 1), 2)
+
+    def test_keeps_at_least_one(self) -> None:
+        # save_freq larger than num_steps yields a single (final) checkpoint.
+        self.assertEqual(resolve_keep_ckpt_count(0.4, 5, 100), 1)
+
+
 if __name__ == "__main__":
     unittest.main()

source/tests/pt/test_training.py

@@ -1180,6 +1180,28 @@ def tearDown(self) -> None:
         os.chdir(self._cwd)
         self._tmpdir.cleanup()
 
+    @TRAINING_TEST_TIMEOUT
+    def test_ckpt_keep_ratio_overrides_keep_counts(self) -> None:
+        config = deepcopy(self.config)
+        config["training"]["ckpt_keep_ratio"] = 0.5
+        trainer = get_trainer(config)
+        # 4 periodic checkpoints; ceil(0.5 * 4) = 2 overrides both the regular
+        # and EMA keep counts.
+        self.assertEqual(trainer.max_ckpt_keep, 2)
+        self.assertEqual(trainer.ema_ckpt_keep, 2)
+        save_ckpt = trainer.save_ckpt
+        ema_save_ckpt = trainer.ema_save_ckpt
+        trainer.run()
+
+        self.assertEqual(
+            sorted(path.name for path in Path(".").glob(f"{save_ckpt}-*.pt")),
+            [f"{save_ckpt}-3.pt", f"{save_ckpt}-4.pt"],
+        )
+        self.assertEqual(
+            sorted(path.name for path in Path(".").glob(f"{ema_save_ckpt}-*.pt")),
+            [f"{ema_save_ckpt}-3.pt", f"{ema_save_ckpt}-4.pt"],
+        )
+
     @TRAINING_TEST_TIMEOUT
     def test_save_dir_redirects_checkpoints_with_local_symlinks(self) -> None:
         config = deepcopy(self.config)