From 636ab5db541403c344143000ae9c09878ed3f09f Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Fri, 9 Oct 2026 13:06:37 +0800
Subject: [PATCH 01/10] fix(train): gate every checkpoint publish with
 NonFiniteGradGuard

Full-validation top-K saves bypassed the non-finite gradient check that
regular/EMA checkpoints already run. Expose one trainer entry that
validates and resets before publication, call it from regular and
validation-best (live + EMA) boundaries on PT and pt_expt, and defer
top-K metadata commit until the write succeeds so a failed boundary
leaves bookkeeping unchanged.

Fixes #5816
---
 deepmd/pt/train/training.py                 | 53 ++++++++++----
 deepmd/pt_expt/train/training.py            | 19 ++++-
 deepmd/pt_expt/train/validation.py          | 75 ++++++++++++++++---
 source/tests/pt/test_validation.py          | 79 +++++++++++++++++++++
 source/tests/pt_expt/test_train_gradient.py | 71 ++++++++++++++++++
 5 files changed, 268 insertions(+), 29 deletions(-)

diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index 7a7210962e..2d0159ffb2 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -1749,22 +1749,14 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
                     step_id=_step_id,
                     display_step=display_step_id,
                     lr=cur_lr,
-                    save_checkpoint=(
-                        self.save_model_merged
-                        if self._lora_enabled
-                        else self.save_model
-                    ),
+                    save_checkpoint=self._save_full_validation_checkpoint,
                 )
             if self.ema_full_validator is not None:
                 self.ema_full_validator.run(
                     step_id=_step_id,
                     display_step=display_step_id,
                     lr=cur_lr,
-                    save_checkpoint=(
-                        self.save_ema_model_merged
-                        if self._lora_enabled
-                        else self.save_ema_model
-                    ),
+                    save_checkpoint=self._save_full_validation_ema_checkpoint,
                 )
 
             should_save_checkpoint = display_step_id == self.num_steps or (
@@ -1773,11 +1765,8 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
                 and _step_id != self.start_step
             )
             if should_save_checkpoint:
-                # Abort before writing if any gradient norm since the previous
-                # checkpoint was non-finite.
-                self.nonfinite_grad_guard.raise_if_nonfinite(
-                    self.wrapper.named_parameters
-                )
+                # One logical publication boundary for live (+ EMA) weights.
+                self.ensure_finite_gradients_for_checkpoint()
             if should_save_checkpoint and (
                 self.sharding.enabled or self.rank == 0 or dist.get_rank() == 0
             ):
@@ -2049,6 +2038,40 @@ def _write_checkpoint(
         torch.save(checkpoint_data, save_path)
         store.prune(save_path)
 
+    def ensure_finite_gradients_for_checkpoint(self) -> None:
+        """Validate and reset non-finite gradient state before publication.
+
+        Every training-side checkpoint boundary — regular, EMA, validation-best,
+        and EMA-validation-best — must call this once before any checkpoint
+        bytes or publication metadata are written. A regular checkpoint that
+        materializes both live and EMA weights is one logical boundary: the
+        check runs before either file is written and is not repeated between
+        them.
+        """
+        self.nonfinite_grad_guard.raise_if_nonfinite(self.wrapper.named_parameters)
+
+    def _save_full_validation_checkpoint(
+        self,
+        save_path: str | Path,
+        lr: float = 0.0,
+        step: int = 0,
+    ) -> None:
+        """Save a live-weight checkpoint selected by full validation."""
+        self.ensure_finite_gradients_for_checkpoint()
+        save = self.save_model_merged if self._lora_enabled else self.save_model
+        save(save_path, lr=lr, step=step)
+
+    def _save_full_validation_ema_checkpoint(
+        self,
+        save_path: str | Path,
+        lr: float = 0.0,
+        step: int = 0,
+    ) -> None:
+        """Save an EMA-weight checkpoint selected by EMA full validation."""
+        self.ensure_finite_gradients_for_checkpoint()
+        save = self.save_ema_model_merged if self._lora_enabled else self.save_ema_model
+        save(save_path, lr=lr, step=step)
+
     def save_model(
         self,
         save_path: str | Path,
diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py
index 9130f5753b..0f3ee80fb3 100644
--- a/deepmd/pt_expt/train/training.py
+++ b/deepmd/pt_expt/train/training.py
@@ -2955,10 +2955,21 @@ def checkpoint_is_collective(self) -> bool:
         """Whether assembling a checkpoint needs every rank."""
         return self.sharding.enabled
 
-    def save_checkpoint(self, step: int) -> None:
-        # Abort before writing if any gradient norm since the previous
-        # checkpoint was non-finite, so a diverged interval is not persisted.
+    def ensure_finite_gradients_for_checkpoint(self) -> None:
+        """Validate and reset non-finite gradient state before publication.
+
+        Every training-side checkpoint boundary — regular, EMA, validation-best,
+        and EMA-validation-best — must call this once before any checkpoint
+        bytes or publication metadata are written. A regular checkpoint that
+        materializes both live and EMA weights is one logical boundary: the
+        check runs before either file is written and is not repeated between
+        them.
+        """
         self.nonfinite_grad_guard.raise_if_nonfinite(self.wrapper.named_parameters)
+
+    def save_checkpoint(self, step: int) -> None:
+        # One logical publication boundary for live (+ EMA) weights.
+        self.ensure_finite_gradients_for_checkpoint()
         ckpt_path = self.ckpt_store.path_for(step)
         self._save_checkpoint_to_path(ckpt_path, step=step)
         if self.rank == 0:
@@ -2980,6 +2991,7 @@ def _save_full_validation_checkpoint(
     ) -> None:
         """Save a checkpoint selected by full validation."""
         del lr
+        self.ensure_finite_gradients_for_checkpoint()
         self._save_checkpoint_to_path(save_path, step=step)
 
     def _save_full_validation_ema_checkpoint(
@@ -2994,6 +3006,7 @@ def _save_full_validation_ema_checkpoint(
         so the shadow has to be applied again while writing it.
         """
         del lr
+        self.ensure_finite_gradients_for_checkpoint()
         self._save_checkpoint_to_path(save_path, step=step, use_ema_weights=True)
 
     def _save_checkpoint_to_path(
diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index 746b6df91a..542ee94f3d 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -289,6 +289,7 @@ def __init__(
             )
 
         self.topk_records = self._load_topk_records()
+        self._pending_topk_records: list[BestCheckpointRecord] | None = None
         self._sync_state_store()
         if self.rank == 0:
             self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
@@ -350,8 +351,14 @@ def run(
                 if (self.is_distributed and self.sharding.enabled) or self.rank == 0:
                     save_checkpoint(Path(save_path[0]), lr=lr, step=step_id)
                 if self.rank == 0:
+                    # Commit top-K bookkeeping only after the checkpoint bytes
+                    # are written, so a failed publication boundary (including
+                    # a non-finite gradient abort) leaves metadata unchanged.
+                    self._commit_pending_best_state()
                     self._reconcile_best_checkpoints()
             except Exception as exc:
+                if self.rank == 0:
+                    self._pending_topk_records = None
                 caught_exception = exc
                 error_message = (
                     "Full validation failed while saving the best checkpoint:\n"
@@ -408,12 +415,21 @@ def _evaluate(self, display_step: int) -> FullValidationResult:
                 f"validation dataset: {self.metric_name.upper()}."
             )
 
-        # === Step 3. Update Best Tracking ===
+        # === Step 3. Propose / commit best tracking ===
+        # When a checkpoint file will be written, top-K metadata stays pending
+        # until that write succeeds (see run()). Tracking-only updates commit
+        # immediately because they never serialize weights.
         selected_metric_value = float(metrics[self.metric_key])
-        saved_best_path = self._update_best_state(
+        proposed_records, saved_best_path = self._propose_best_checkpoint(
             display_step=display_step,
             selected_metric_value=selected_metric_value,
         )
+        self._pending_topk_records = None
+        if saved_best_path is not None:
+            self._pending_topk_records = proposed_records
+        elif proposed_records is not None:
+            self.topk_records = proposed_records
+            self._sync_state_store()
         return FullValidationResult(
             display_step=display_step,
             metrics=metrics,
@@ -696,13 +712,21 @@ def predict_batch(
             prediction["virial"] = batch_prediction["virial"]
         return prediction
 
-    def _update_best_state(
+    def _propose_best_checkpoint(
         self,
         *,
         display_step: int,
         selected_metric_value: float,
-    ) -> str | None:
-        """Update the top-K records and return the checkpoint path to save."""
+    ) -> tuple[list[BestCheckpointRecord] | None, str | None]:
+        """Compute the next top-K set and optional save path without committing.
+
+        Returns
+        -------
+        tuple[list[BestCheckpointRecord] | None, str | None]
+            ``(updated_records, save_path)``. ``updated_records`` is ``None``
+            when the candidate misses the top-K window. ``save_path`` is set
+            only when ``save_best`` is enabled and the candidate is retained.
+        """
         candidate = BestCheckpointRecord(
             metric=selected_metric_value,
             step=display_step,
@@ -714,14 +738,43 @@ def _update_best_state(
         updated_records.sort()
         updated_records = updated_records[: self.max_best_ckpt]
         if candidate not in updated_records:
-            return None
+            return None, None
+        if not self.save_best:
+            return updated_records, None
+        candidate_rank = updated_records.index(candidate) + 1
+        return updated_records, str(
+            self._best_checkpoint_path(display_step, candidate_rank)
+        )
 
-        self.topk_records = updated_records
+    def _commit_pending_best_state(self) -> None:
+        """Persist a deferred top-K update after a successful checkpoint write."""
+        if self._pending_topk_records is None:
+            return
+        self.topk_records = self._pending_topk_records
         self._sync_state_store()
-        if not self.save_best:
-            return None
-        candidate_rank = self.topk_records.index(candidate) + 1
-        return str(self._best_checkpoint_path(display_step, candidate_rank))
+        self._pending_topk_records = None
+
+    def _update_best_state(
+        self,
+        *,
+        display_step: int,
+        selected_metric_value: float,
+    ) -> str | None:
+        """Update the top-K records and return the checkpoint path to save.
+
+        This helper commits immediately and is used by unit tests and callers
+        that manage serialization themselves. The training ``run()`` path uses
+        :meth:`_propose_best_checkpoint` so metadata can wait on a successful
+        write.
+        """
+        updated_records, save_path = self._propose_best_checkpoint(
+            display_step=display_step,
+            selected_metric_value=selected_metric_value,
+        )
+        if updated_records is not None:
+            self.topk_records = updated_records
+            self._sync_state_store()
+        return save_path
 
     def _sync_state_store(self) -> None:
         """Synchronize top-K validation state into the configured state store."""
diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py
index 767b9967e4..9899bf093d 100644
--- a/source/tests/pt/test_validation.py
+++ b/source/tests/pt/test_validation.py
@@ -990,3 +990,82 @@ def test_predict_outputs_emits_real_and_magnetic_forces(self) -> None:
         np.testing.assert_array_equal(
             prediction["mask_mag"].astype(bool), expected_mask
         )
+
+
+class TestFullValidatorCheckpointGate(unittest.TestCase):
+    def _make_validator(self, tmpdir: str) -> FullValidator:
+        return FullValidator(
+            validating_params={
+                "full_validation": True,
+                "validation_freq": 1,
+                "save_best": True,
+                "max_best_ckpt": 1,
+                "validation_metric": "E:MAE",
+                "full_val_file": str(Path(tmpdir) / "val.log"),
+                "full_val_start": 0.0,
+            },
+            validation_data=_DummyValidationData(),
+            model=_DummyModel(),
+            state_store={},
+            num_steps=10,
+            rank=0,
+            restart_training=False,
+            checkpoint_dir=Path(tmpdir),
+        )
+
+    def test_full_validator_defers_topk_until_successful_save(self) -> None:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = self._make_validator(tmpdir)
+            writes: list[Path] = []
+
+            def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
+                del lr, step
+                Path(path).write_text("ok")
+                writes.append(Path(path))
+
+            with patch.object(
+                validator,
+                "evaluate_all_systems",
+                return_value={validator.metric_key: 0.5},
+            ):
+                result = validator.run(
+                    step_id=1,
+                    display_step=1,
+                    lr=1e-3,
+                    save_checkpoint=save_checkpoint,
+                )
+
+            self.assertIsNotNone(result)
+            self.assertEqual(len(writes), 1)
+            self.assertTrue(writes[0].exists())
+            self.assertEqual(len(validator.topk_records), 1)
+            self.assertEqual(validator.topk_records[0].step, 1)
+            self.assertIsNone(validator._pending_topk_records)
+
+    def test_full_validator_discards_topk_when_save_raises(self) -> None:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = self._make_validator(tmpdir)
+
+            def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
+                del path, lr, step
+                raise RuntimeError("Non-finite gradient norm; training has diverged.")
+
+            with patch.object(
+                validator,
+                "evaluate_all_systems",
+                return_value={validator.metric_key: 0.25},
+            ):
+                with self.assertRaisesRegex(RuntimeError, "Non-finite"):
+                    validator.run(
+                        step_id=2,
+                        display_step=2,
+                        lr=1e-3,
+                        save_checkpoint=save_checkpoint,
+                    )
+
+            self.assertEqual(validator.topk_records, [])
+            self.assertIsNone(validator._pending_topk_records)
+            self.assertEqual(
+                list(Path(tmpdir).glob("best.ckpt-*.pt")),
+                [],
+            )
diff --git a/source/tests/pt_expt/test_train_gradient.py b/source/tests/pt_expt/test_train_gradient.py
index f20912f57b..b9d703f4f5 100644
--- a/source/tests/pt_expt/test_train_gradient.py
+++ b/source/tests/pt_expt/test_train_gradient.py
@@ -111,5 +111,76 @@ def test_resets_after_check(self) -> None:
         guard.raise_if_nonfinite(self._named(1.0))
 
 
+class TestCheckpointPublicationGuard(unittest.TestCase):
+    """Publication-boundary contract for NonFiniteGradGuard (#5816)."""
+
+    @staticmethod
+    def _named(grad_value: float = 1.0):
+        p = torch.nn.Parameter(torch.zeros(2, device="cpu"))
+        p.grad = torch.full((2,), grad_value, device="cpu")
+        return lambda: [("layer.weight", p)]
+
+    def test_validation_best_raises_before_serialize_and_keeps_metadata(self) -> None:
+        # A non-finite norm remains sticky after a later finite update, so a
+        # validation-best publish before the regular interval must still abort.
+        guard = NonFiniteGradGuard()
+        guard.update(torch.tensor(float("nan"), device="cpu"))
+        guard.update(torch.tensor(1.0, device="cpu"))
+
+        writes: list[str] = []
+        topk_committed = {"value": False}
+
+        def ensure() -> None:
+            guard.raise_if_nonfinite(self._named(float("nan")))
+
+        def save_validation_best() -> None:
+            ensure()
+            writes.append("validation-best")
+            topk_committed["value"] = True
+
+        with self.assertRaises(RuntimeError):
+            save_validation_best()
+        self.assertEqual(writes, [])
+        self.assertFalse(topk_committed["value"])
+
+    def test_ema_validation_best_raises_before_serialize(self) -> None:
+        guard = NonFiniteGradGuard()
+        guard.update(torch.tensor(float("inf"), device="cpu"))
+        writes: list[str] = []
+
+        def save_ema_validation_best() -> None:
+            guard.raise_if_nonfinite(self._named(1.0))
+            writes.append("ema-validation-best")
+
+        with self.assertRaises(RuntimeError):
+            save_ema_validation_best()
+        self.assertEqual(writes, [])
+
+    def test_live_and_ema_regular_boundary_validates_once(self) -> None:
+        guard = NonFiniteGradGuard()
+        guard.update(torch.tensor(1.0, device="cpu"))
+        checks = {"count": 0}
+        writes: list[str] = []
+
+        real_raise = guard.raise_if_nonfinite
+
+        def counting_raise(named_parameters) -> None:
+            checks["count"] += 1
+            real_raise(named_parameters)
+
+        guard.raise_if_nonfinite = counting_raise  # type: ignore[method-assign]
+
+        # One logical boundary: gate once, then write live and EMA.
+        guard.raise_if_nonfinite(self._named(1.0))
+        writes.append("live")
+        writes.append("ema")
+
+        self.assertEqual(checks["count"], 1)
+        self.assertEqual(writes, ["live", "ema"])
+        # The successful boundary cleared the flag; a second gate is a no-op.
+        guard.raise_if_nonfinite(self._named(1.0))
+        self.assertEqual(checks["count"], 2)
+
+
 if __name__ == "__main__":
     unittest.main()

From 0fbbfbb8ae2c2b28ba7701864fafceba30fcf80c Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Fri, 9 Oct 2026 13:35:54 +0800
Subject: [PATCH 02/10] fix(train): all-rank full-val gate + lock publish paths
 in tests

Full-validation checkpoint publication now enters the save callback on
every rank so NonFiniteGradGuard validates and resets consistently,
matching PT regular saves; non-chief writers still no-op inside
_write_checkpoint. Replace stub publication-boundary tests with thin
Trainer.__new__ coverage of the real save entry points, and assert
distributed non-chief ranks enter the save callback.
---
 deepmd/pt_expt/train/validation.py          |   9 +-
 source/tests/pt/test_validation.py          |  47 +++++++++
 source/tests/pt_expt/test_train_gradient.py | 111 ++++++++++++--------
 3 files changed, 118 insertions(+), 49 deletions(-)

diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index 542ee94f3d..bf9f674a7a 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -346,10 +346,11 @@ def run(
 
         if save_path[0] is not None:
             try:
-                # Assembling a checkpoint from shards is collective, so every
-                # rank enters it once any training state is sharded.
-                if (self.is_distributed and self.sharding.enabled) or self.rank == 0:
-                    save_checkpoint(Path(save_path[0]), lr=lr, step=step_id)
+                # Every rank enters the publication boundary so NonFiniteGradGuard
+                # validates and resets consistently (mirroring PT regular saves).
+                # Writers and sharded collectives stay inside the save callback /
+                # Trainer._write_checkpoint; non-chief ranks return before disk I/O.
+                save_checkpoint(Path(save_path[0]), lr=lr, step=step_id)
                 if self.rank == 0:
                     # Commit top-K bookkeeping only after the checkpoint bytes
                     # are written, so a failed publication boundary (including
diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py
index 9899bf093d..0931cf6a2e 100644
--- a/source/tests/pt/test_validation.py
+++ b/source/tests/pt/test_validation.py
@@ -1069,3 +1069,50 @@ def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
                 list(Path(tmpdir).glob("best.ckpt-*.pt")),
                 [],
             )
+
+    def test_distributed_non_chief_enters_save_callback(self) -> None:
+        """Non-chief ranks must enter save so the non-finite gate can reset."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = self._make_validator(tmpdir)
+            validator.rank = 1
+            validator.is_distributed = True
+            writes: list[Path] = []
+
+            def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
+                del lr, step
+                Path(path).write_text("ok")
+                writes.append(Path(path))
+
+            def broadcast_save_path(holder: list, src: int = 0) -> None:
+                del src
+                holder[0] = str(Path(tmpdir) / "best.ckpt-7-1.pt")
+
+            with (
+                patch(
+                    "deepmd.pt_expt.train.validation.dist.barrier",
+                ),
+                patch(
+                    "deepmd.pt_expt.train.validation.dist.broadcast_object_list",
+                    side_effect=broadcast_save_path,
+                ),
+                patch(
+                    "deepmd.pt_expt.train.validation.dist.get_world_size",
+                    return_value=2,
+                ),
+                patch(
+                    "deepmd.pt_expt.train.validation.dist.all_gather_object",
+                    side_effect=lambda out, obj: out.__setitem__(
+                        slice(None), [obj, None]
+                    ),
+                ),
+            ):
+                result = validator.run(
+                    step_id=7,
+                    display_step=7,
+                    lr=1e-3,
+                    save_checkpoint=save_checkpoint,
+                )
+
+            self.assertIsNone(result)
+            self.assertEqual(len(writes), 1)
+            self.assertTrue(writes[0].exists())
diff --git a/source/tests/pt_expt/test_train_gradient.py b/source/tests/pt_expt/test_train_gradient.py
index b9d703f4f5..966eedc415 100644
--- a/source/tests/pt_expt/test_train_gradient.py
+++ b/source/tests/pt_expt/test_train_gradient.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import unittest
+from pathlib import Path
 
 import torch
 
@@ -112,73 +113,93 @@ def test_resets_after_check(self) -> None:
 
 
 class TestCheckpointPublicationGuard(unittest.TestCase):
-    """Publication-boundary contract for NonFiniteGradGuard (#5816)."""
+    """Publication-boundary contract via real Trainer save entry points (#5816)."""
 
     @staticmethod
-    def _named(grad_value: float = 1.0):
+    def _named_parameters(grad_value: float = 1.0):
         p = torch.nn.Parameter(torch.zeros(2, device="cpu"))
         p.grad = torch.full((2,), grad_value, device="cpu")
         return lambda: [("layer.weight", p)]
 
-    def test_validation_best_raises_before_serialize_and_keeps_metadata(self) -> None:
-        # A non-finite norm remains sticky after a later finite update, so a
-        # validation-best publish before the regular interval must still abort.
-        guard = NonFiniteGradGuard()
-        guard.update(torch.tensor(float("nan"), device="cpu"))
-        guard.update(torch.tensor(1.0, device="cpu"))
-
-        writes: list[str] = []
-        topk_committed = {"value": False}
-
-        def ensure() -> None:
-            guard.raise_if_nonfinite(self._named(float("nan")))
-
-        def save_validation_best() -> None:
-            ensure()
-            writes.append("validation-best")
-            topk_committed["value"] = True
-
+    def _bare_trainer(self, *, with_ema: bool = False):
+        from types import SimpleNamespace
+
+        from deepmd.pt_expt.train.training import Trainer
+
+        trainer = Trainer.__new__(Trainer)
+        trainer.nonfinite_grad_guard = NonFiniteGradGuard()
+        trainer.wrapper = SimpleNamespace(named_parameters=self._named_parameters(1.0))
+        trainer.rank = 0
+        trainer.model_ema = object() if with_ema else None
+        trainer.ckpt_store = SimpleNamespace(
+            path_for=lambda step: Path(f"model.ckpt-{step}.pt"),
+            publish=lambda path: None,
+            prune=lambda path: None,
+        )
+        trainer.ema_ckpt_store = SimpleNamespace(
+            path_for=lambda step: Path(f"model.ckpt-ema-{step}.pt"),
+            publish=lambda path: None,
+            prune=lambda path: None,
+        )
+        return trainer
+
+    def test_validation_best_raises_before_serialize(self) -> None:
+        # Sticky non-finite after a later finite update must still abort
+        # validation-best publication before any serialize callback runs.
+        trainer = self._bare_trainer()
+        trainer.nonfinite_grad_guard.update(torch.tensor(float("nan"), device="cpu"))
+        trainer.nonfinite_grad_guard.update(torch.tensor(1.0, device="cpu"))
+        writes: list[object] = []
+
+        def fake_write(path, *, step: int, use_ema_weights: bool = False) -> None:
+            del step, use_ema_weights
+            writes.append(path)
+
+        trainer._save_checkpoint_to_path = fake_write  # type: ignore[method-assign]
         with self.assertRaises(RuntimeError):
-            save_validation_best()
+            trainer._save_full_validation_checkpoint(Path("best.pt"), step=3)
         self.assertEqual(writes, [])
-        self.assertFalse(topk_committed["value"])
 
     def test_ema_validation_best_raises_before_serialize(self) -> None:
-        guard = NonFiniteGradGuard()
-        guard.update(torch.tensor(float("inf"), device="cpu"))
-        writes: list[str] = []
+        trainer = self._bare_trainer()
+        trainer.nonfinite_grad_guard.update(torch.tensor(float("inf"), device="cpu"))
+        writes: list[object] = []
 
-        def save_ema_validation_best() -> None:
-            guard.raise_if_nonfinite(self._named(1.0))
-            writes.append("ema-validation-best")
+        def fake_write(path, *, step: int, use_ema_weights: bool = False) -> None:
+            del step
+            writes.append((path, use_ema_weights))
 
+        trainer._save_checkpoint_to_path = fake_write  # type: ignore[method-assign]
         with self.assertRaises(RuntimeError):
-            save_ema_validation_best()
+            trainer._save_full_validation_ema_checkpoint(Path("best-ema.pt"), step=4)
         self.assertEqual(writes, [])
 
-    def test_live_and_ema_regular_boundary_validates_once(self) -> None:
-        guard = NonFiniteGradGuard()
-        guard.update(torch.tensor(1.0, device="cpu"))
+    def test_regular_boundary_validates_once_for_live_and_ema(self) -> None:
+        trainer = self._bare_trainer(with_ema=True)
+        trainer.nonfinite_grad_guard.update(torch.tensor(1.0, device="cpu"))
         checks = {"count": 0}
-        writes: list[str] = []
+        writes: list[bool] = []
+        real_ensure = trainer.ensure_finite_gradients_for_checkpoint
 
-        real_raise = guard.raise_if_nonfinite
-
-        def counting_raise(named_parameters) -> None:
+        def counting_ensure() -> None:
             checks["count"] += 1
-            real_raise(named_parameters)
+            real_ensure()
 
-        guard.raise_if_nonfinite = counting_raise  # type: ignore[method-assign]
+        def fake_write(path, *, step: int, use_ema_weights: bool = False) -> None:
+            del path, step
+            writes.append(use_ema_weights)
 
-        # One logical boundary: gate once, then write live and EMA.
-        guard.raise_if_nonfinite(self._named(1.0))
-        writes.append("live")
-        writes.append("ema")
+        trainer.ensure_finite_gradients_for_checkpoint = (  # type: ignore[method-assign]
+            counting_ensure
+        )
+        trainer._save_checkpoint_to_path = fake_write  # type: ignore[method-assign]
+
+        trainer.save_checkpoint(5)
 
         self.assertEqual(checks["count"], 1)
-        self.assertEqual(writes, ["live", "ema"])
-        # The successful boundary cleared the flag; a second gate is a no-op.
-        guard.raise_if_nonfinite(self._named(1.0))
+        self.assertEqual(writes, [False, True])
+        # Successful boundary cleared the flag; a later gate is a no-op.
+        trainer.ensure_finite_gradients_for_checkpoint()
         self.assertEqual(checks["count"], 2)
 
 

From 9e81f6945e13bf56078ffb44e8094d1aae8d986e Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 9 Oct 2026 05:37:00 +0000
Subject: [PATCH 03/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/pt_expt/test_train_gradient.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/source/tests/pt_expt/test_train_gradient.py b/source/tests/pt_expt/test_train_gradient.py
index 966eedc415..8434cdde6e 100644
--- a/source/tests/pt_expt/test_train_gradient.py
+++ b/source/tests/pt_expt/test_train_gradient.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import unittest
-from pathlib import Path
+from pathlib import (
+    Path,
+)
 
 import torch
 
@@ -122,9 +124,13 @@ def _named_parameters(grad_value: float = 1.0):
         return lambda: [("layer.weight", p)]
 
     def _bare_trainer(self, *, with_ema: bool = False):
-        from types import SimpleNamespace
+        from types import (
+            SimpleNamespace,
+        )
 
-        from deepmd.pt_expt.train.training import Trainer
+        from deepmd.pt_expt.train.training import (
+            Trainer,
+        )
 
         trainer = Trainer.__new__(Trainer)
         trainer.nonfinite_grad_guard = NonFiniteGradGuard()

From 4e1f65020acf4894617bc9704f1275a84bf944b8 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Fri, 9 Oct 2026 17:33:01 +0800
Subject: [PATCH 04/10] fix(train): stage top-K into checkpoint before
 publication
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Deferred commit after save left train_infos one validation behind, so
restart reconciliation could delete the just-written best file. Stage
proposed records into state_store before serialization, keep a rollback
snapshot for failed boundaries, and cover save→reload reconcile.
---
 deepmd/pt_expt/train/validation.py |  43 +++++++--
 source/tests/pt/test_validation.py | 144 +++++++++++++++++++++++++++++
 2 files changed, 178 insertions(+), 9 deletions(-)

diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index bf9f674a7a..5d18407f06 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -290,6 +290,7 @@ def __init__(
 
         self.topk_records = self._load_topk_records()
         self._pending_topk_records: list[BestCheckpointRecord] | None = None
+        self._rollback_topk_records: list[BestCheckpointRecord] | None = None
         self._sync_state_store()
         if self.rank == 0:
             self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
@@ -346,20 +347,23 @@ def run(
 
         if save_path[0] is not None:
             try:
+                # Stage proposed top-K into the live state_store before
+                # serialization so the written checkpoint carries the records
+                # that name this file. Keep a rollback snapshot until the write
+                # either commits or aborts (non-finite gate, I/O, etc.).
+                if self.rank == 0:
+                    self._stage_pending_best_state()
                 # Every rank enters the publication boundary so NonFiniteGradGuard
                 # validates and resets consistently (mirroring PT regular saves).
                 # Writers and sharded collectives stay inside the save callback /
                 # Trainer._write_checkpoint; non-chief ranks return before disk I/O.
                 save_checkpoint(Path(save_path[0]), lr=lr, step=step_id)
                 if self.rank == 0:
-                    # Commit top-K bookkeeping only after the checkpoint bytes
-                    # are written, so a failed publication boundary (including
-                    # a non-finite gradient abort) leaves metadata unchanged.
                     self._commit_pending_best_state()
                     self._reconcile_best_checkpoints()
             except Exception as exc:
                 if self.rank == 0:
-                    self._pending_topk_records = None
+                    self._rollback_pending_best_state()
                 caught_exception = exc
                 error_message = (
                     "Full validation failed while saving the best checkpoint:\n"
@@ -418,7 +422,8 @@ def _evaluate(self, display_step: int) -> FullValidationResult:
 
         # === Step 3. Propose / commit best tracking ===
         # When a checkpoint file will be written, top-K metadata stays pending
-        # until that write succeeds (see run()). Tracking-only updates commit
+        # until run() stages it into state_store for serialization and then
+        # commits or rolls back after the write. Tracking-only updates commit
         # immediately because they never serialize weights.
         selected_metric_value = float(metrics[self.metric_key])
         proposed_records, saved_best_path = self._propose_best_checkpoint(
@@ -747,13 +752,33 @@ def _propose_best_checkpoint(
             self._best_checkpoint_path(display_step, candidate_rank)
         )
 
-    def _commit_pending_best_state(self) -> None:
-        """Persist a deferred top-K update after a successful checkpoint write."""
+    def _stage_pending_best_state(self) -> None:
+        """Expose proposed top-K in ``state_store`` before checkpoint serialization.
+
+        The trainer serializes the validator's ``state_store`` (via wrapper
+        ``train_infos`` / EMA validation state) inside the save callback. The
+        proposed records must therefore be live before that call, while a
+        rollback snapshot preserves the previous bookkeeping until the write
+        commits or fails.
+        """
         if self._pending_topk_records is None:
             return
+        self._rollback_topk_records = list(self.topk_records)
         self.topk_records = self._pending_topk_records
         self._sync_state_store()
+
+    def _commit_pending_best_state(self) -> None:
+        """Keep staged top-K after a successful checkpoint write."""
+        self._pending_topk_records = None
+        self._rollback_topk_records = None
+
+    def _rollback_pending_best_state(self) -> None:
+        """Restore pre-proposal top-K after a failed publication boundary."""
+        if self._rollback_topk_records is not None:
+            self.topk_records = self._rollback_topk_records
+            self._sync_state_store()
         self._pending_topk_records = None
+        self._rollback_topk_records = None
 
     def _update_best_state(
         self,
@@ -765,8 +790,8 @@ def _update_best_state(
 
         This helper commits immediately and is used by unit tests and callers
         that manage serialization themselves. The training ``run()`` path uses
-        :meth:`_propose_best_checkpoint` so metadata can wait on a successful
-        write.
+        :meth:`_propose_best_checkpoint` so metadata can be staged for the
+        write and rolled back if publication fails.
         """
         updated_records, save_path = self._propose_best_checkpoint(
             display_step=display_step,
diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py
index 0931cf6a2e..df0a244c7e 100644
--- a/source/tests/pt/test_validation.py
+++ b/source/tests/pt/test_validation.py
@@ -1017,9 +1017,14 @@ def test_full_validator_defers_topk_until_successful_save(self) -> None:
         with tempfile.TemporaryDirectory() as tmpdir:
             validator = self._make_validator(tmpdir)
             writes: list[Path] = []
+            serialized_topk: list[list[dict]] = []
 
             def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
                 del lr, step
+                # The checkpoint bytes must already see the proposed top-K.
+                serialized_topk.append(
+                    list(validator.state_store[validator.topk_records_info_key])
+                )
                 Path(path).write_text("ok")
                 writes.append(Path(path))
 
@@ -1041,6 +1046,8 @@ def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
             self.assertEqual(len(validator.topk_records), 1)
             self.assertEqual(validator.topk_records[0].step, 1)
             self.assertIsNone(validator._pending_topk_records)
+            self.assertIsNone(validator._rollback_topk_records)
+            self.assertEqual(serialized_topk, [[{"metric": 0.5, "step": 1}]])
 
     def test_full_validator_discards_topk_when_save_raises(self) -> None:
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -1065,11 +1072,148 @@ def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
 
             self.assertEqual(validator.topk_records, [])
             self.assertIsNone(validator._pending_topk_records)
+            self.assertIsNone(validator._rollback_topk_records)
+            self.assertEqual(
+                validator.state_store.get(validator.topk_records_info_key, []),
+                [],
+            )
             self.assertEqual(
                 list(Path(tmpdir).glob("best.ckpt-*.pt")),
                 [],
             )
 
+    def test_full_validator_serialized_topk_survives_restart_reconcile(self) -> None:
+        """Improved best ckpt must serialize its own record so restart keeps it."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            state_store: dict = {}
+            validator = FullValidator(
+                validating_params={
+                    "full_validation": True,
+                    "validation_freq": 1,
+                    "save_best": True,
+                    "max_best_ckpt": 1,
+                    "validation_metric": "E:MAE",
+                    "full_val_file": str(Path(tmpdir) / "val.log"),
+                    "full_val_start": 0.0,
+                },
+                validation_data=_DummyValidationData(),
+                model=_DummyModel(),
+                state_store=state_store,
+                num_steps=10,
+                rank=0,
+                restart_training=False,
+                checkpoint_dir=Path(tmpdir),
+            )
+            serialized_stores: list[dict] = []
+
+            def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
+                del lr, step
+                # Snapshot what a real torch.save would capture from train_infos.
+                serialized_stores.append(
+                    {
+                        validator.metric_name_info_key: state_store[
+                            validator.metric_name_info_key
+                        ],
+                        validator.topk_records_info_key: [
+                            dict(record)
+                            for record in state_store[validator.topk_records_info_key]
+                        ],
+                    }
+                )
+                Path(path).write_text(f"ckpt-{path.name}")
+                # Leave older best files in place; reconcile decides retention.
+
+            metrics = {1: 0.5, 2: 0.1}
+            for display_step, metric in metrics.items():
+                with patch.object(
+                    validator,
+                    "evaluate_all_systems",
+                    return_value={validator.metric_key: metric},
+                ):
+                    validator.run(
+                        step_id=display_step,
+                        display_step=display_step,
+                        lr=1e-3,
+                        save_checkpoint=save_checkpoint,
+                    )
+
+            self.assertEqual(len(serialized_stores), 2)
+            self.assertEqual(
+                serialized_stores[1][validator.topk_records_info_key],
+                [{"metric": 0.1, "step": 2}],
+            )
+            # Restart from the improved checkpoint's serialized train_infos.
+            restarted = FullValidator(
+                validating_params={
+                    "full_validation": True,
+                    "validation_freq": 1,
+                    "save_best": True,
+                    "max_best_ckpt": 1,
+                    "validation_metric": "E:MAE",
+                    "full_val_file": str(Path(tmpdir) / "val.log"),
+                    "full_val_start": 0.0,
+                },
+                validation_data=_DummyValidationData(),
+                model=_DummyModel(),
+                state_store=dict(serialized_stores[1]),
+                num_steps=10,
+                rank=0,
+                restart_training=True,
+                checkpoint_dir=Path(tmpdir),
+            )
+            self.assertEqual(len(restarted.topk_records), 1)
+            self.assertEqual(restarted.topk_records[0].step, 2)
+            remaining = sorted(Path(tmpdir).glob("best.ckpt-*.pt"))
+            self.assertEqual(len(remaining), 1)
+            self.assertIn("best.ckpt-2", remaining[0].name)
+
+    def test_full_validator_rolls_back_previous_topk_on_failed_save(self) -> None:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = self._make_validator(tmpdir)
+
+            def save_ok(path: Path, lr: float = 0.0, step: int = 0) -> None:
+                del lr, step
+                Path(path).write_text("ok")
+
+            def save_fail(path: Path, lr: float = 0.0, step: int = 0) -> None:
+                del path, lr, step
+                raise RuntimeError("Non-finite gradient norm; training has diverged.")
+
+            with patch.object(
+                validator,
+                "evaluate_all_systems",
+                return_value={validator.metric_key: 0.5},
+            ):
+                validator.run(
+                    step_id=1,
+                    display_step=1,
+                    lr=1e-3,
+                    save_checkpoint=save_ok,
+                )
+            self.assertEqual(validator.topk_records[0].step, 1)
+
+            with patch.object(
+                validator,
+                "evaluate_all_systems",
+                return_value={validator.metric_key: 0.1},
+            ):
+                with self.assertRaisesRegex(RuntimeError, "Non-finite"):
+                    validator.run(
+                        step_id=2,
+                        display_step=2,
+                        lr=1e-3,
+                        save_checkpoint=save_fail,
+                    )
+
+            self.assertEqual(len(validator.topk_records), 1)
+            self.assertEqual(validator.topk_records[0].step, 1)
+            self.assertEqual(
+                validator.state_store[validator.topk_records_info_key],
+                [{"metric": 0.5, "step": 1}],
+            )
+            self.assertIsNone(validator._pending_topk_records)
+            self.assertIsNone(validator._rollback_topk_records)
+
     def test_distributed_non_chief_enters_save_callback(self) -> None:
         """Non-chief ranks must enter save so the non-finite gate can reset."""
         with tempfile.TemporaryDirectory() as tmpdir:

From 0539cd3910c96d5d21d46ecb9eb1dd9e6efaa178 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Fri, 9 Oct 2026 17:46:19 +0800
Subject: [PATCH 05/10] fix(train): commit top-K only after all ranks accept
 save

Rank 0 previously committed and reconciled best-checkpoint metadata
immediately after its local save returned. A non-chief abort then
surfaced via _raise_if_distributed_error, but the file and top-K
record stayed published. Defer commit until the distributed sync
succeeds; on remote failure after a local write, roll back staged
top-K and reconcile away the orphan.
---
 deepmd/pt_expt/train/validation.py | 24 +++++++++---
 source/tests/pt/test_validation.py | 62 ++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 6 deletions(-)

diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index 5d18407f06..6c75114d66 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -346,11 +346,15 @@ def run(
             dist.broadcast_object_list(save_path, src=0)
 
         if save_path[0] is not None:
+            # Stage → save on every rank → sync outcomes → only then commit.
+            # Committing before ``_raise_if_distributed_error`` would let rank 0
+            # publish top-K (and keep the file) when a non-chief rank aborted
+            # the publication boundary.
             try:
                 # Stage proposed top-K into the live state_store before
                 # serialization so the written checkpoint carries the records
-                # that name this file. Keep a rollback snapshot until the write
-                # either commits or aborts (non-finite gate, I/O, etc.).
+                # that name this file. Keep a rollback snapshot until every
+                # rank reports a successful boundary.
                 if self.rank == 0:
                     self._stage_pending_best_state()
                 # Every rank enters the publication boundary so NonFiniteGradGuard
@@ -358,9 +362,6 @@ def run(
                 # Writers and sharded collectives stay inside the save callback /
                 # Trainer._write_checkpoint; non-chief ranks return before disk I/O.
                 save_checkpoint(Path(save_path[0]), lr=lr, step=step_id)
-                if self.rank == 0:
-                    self._commit_pending_best_state()
-                    self._reconcile_best_checkpoints()
             except Exception as exc:
                 if self.rank == 0:
                     self._rollback_pending_best_state()
@@ -373,7 +374,18 @@ def run(
                 error_message = None
                 caught_exception = None
 
-            self._raise_if_distributed_error(error_message, caught_exception)
+            try:
+                self._raise_if_distributed_error(error_message, caught_exception)
+            except Exception:
+                # Local save succeeded but another rank failed: undo staged
+                # top-K and drop the orphaned file before propagating.
+                if self.rank == 0 and caught_exception is None:
+                    self._rollback_pending_best_state()
+                    self._reconcile_best_checkpoints()
+                raise
+            if self.rank == 0:
+                self._commit_pending_best_state()
+                self._reconcile_best_checkpoints()
 
         if self.rank == 0:
             try:
diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py
index df0a244c7e..d85360b2b5 100644
--- a/source/tests/pt/test_validation.py
+++ b/source/tests/pt/test_validation.py
@@ -1214,6 +1214,68 @@ def save_fail(path: Path, lr: float = 0.0, step: int = 0) -> None:
             self.assertIsNone(validator._pending_topk_records)
             self.assertIsNone(validator._rollback_topk_records)
 
+    def test_distributed_remote_save_failure_rolls_back_committed_topk(self) -> None:
+        """Rank 0 must not keep top-K when another rank aborts the save."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = self._make_validator(tmpdir)
+            validator.is_distributed = True
+            writes: list[Path] = []
+
+            def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
+                del lr, step
+                Path(path).write_text("ok")
+                writes.append(Path(path))
+
+            gather_calls = {"n": 0}
+
+            def gather_remote_error(out: list, obj: object) -> None:
+                # Succeed on the post-evaluate sync; fail on the post-save sync.
+                gather_calls["n"] += 1
+                if gather_calls["n"] == 2:
+                    out[:] = [
+                        obj,
+                        "Full validation failed while saving the best checkpoint:\nremote",
+                    ]
+                else:
+                    out[:] = [obj, None]
+
+            with (
+                patch.object(
+                    validator,
+                    "evaluate_all_systems",
+                    return_value={validator.metric_key: 0.5},
+                ),
+                patch("deepmd.pt_expt.train.validation.dist.barrier"),
+                patch(
+                    "deepmd.pt_expt.train.validation.dist.broadcast_object_list",
+                ),
+                patch(
+                    "deepmd.pt_expt.train.validation.dist.get_world_size",
+                    return_value=2,
+                ),
+                patch(
+                    "deepmd.pt_expt.train.validation.dist.all_gather_object",
+                    side_effect=gather_remote_error,
+                ),
+            ):
+                with self.assertRaisesRegex(RuntimeError, "remote"):
+                    validator.run(
+                        step_id=3,
+                        display_step=3,
+                        lr=1e-3,
+                        save_checkpoint=save_checkpoint,
+                    )
+
+            self.assertEqual(len(writes), 1)
+            self.assertEqual(validator.topk_records, [])
+            self.assertIsNone(validator._pending_topk_records)
+            self.assertIsNone(validator._rollback_topk_records)
+            self.assertEqual(
+                validator.state_store.get(validator.topk_records_info_key, []),
+                [],
+            )
+            self.assertEqual(list(Path(tmpdir).glob("best.ckpt-*.pt")), [])
+
     def test_distributed_non_chief_enters_save_callback(self) -> None:
         """Non-chief ranks must enter save so the non-finite gate can reset."""
         with tempfile.TemporaryDirectory() as tmpdir:

From af50c3f6f212304720ddd68986223bae43a81238 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Fri, 9 Oct 2026 17:51:25 +0800
Subject: [PATCH 06/10] fix(train): reconcile orphans after local
 write-then-fail

Remote abort after a successful rank-0 write already rolls back top-K and
reconciles away the orphan. A local exception after bytes land on disk only
rolled back metadata, leaving best.ckpt-* behind. Call reconcile after the
local rollback so both failure paths drop aborted candidates.
---
 deepmd/pt_expt/train/validation.py |  4 +++
 source/tests/pt/test_validation.py | 51 ++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index 6c75114d66..9441b95253 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -364,7 +364,11 @@ def run(
                 save_checkpoint(Path(save_path[0]), lr=lr, step=step_id)
             except Exception as exc:
                 if self.rank == 0:
+                    # Match remote-failure cleanup: restore top-K, then drop any
+                    # bytes already written for the aborted candidate (write-
+                    # then-fail leaves an orphan that rollback alone keeps).
                     self._rollback_pending_best_state()
+                    self._reconcile_best_checkpoints()
                 caught_exception = exc
                 error_message = (
                     "Full validation failed while saving the best checkpoint:\n"
diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py
index d85360b2b5..cd404a0cf8 100644
--- a/source/tests/pt/test_validation.py
+++ b/source/tests/pt/test_validation.py
@@ -1214,6 +1214,57 @@ def save_fail(path: Path, lr: float = 0.0, step: int = 0) -> None:
             self.assertIsNone(validator._pending_topk_records)
             self.assertIsNone(validator._rollback_topk_records)
 
+    def test_local_write_then_fail_reconciles_orphan(self) -> None:
+        """Local save that writes then raises must delete the orphan file."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = self._make_validator(tmpdir)
+
+            def save_ok(path: Path, lr: float = 0.0, step: int = 0) -> None:
+                del lr, step
+                Path(path).write_text("ok")
+
+            def save_write_then_fail(
+                path: Path, lr: float = 0.0, step: int = 0
+            ) -> None:
+                del lr, step
+                Path(path).write_text("orphan")
+                raise RuntimeError("post-write failure")
+
+            with patch.object(
+                validator,
+                "evaluate_all_systems",
+                return_value={validator.metric_key: 0.5},
+            ):
+                validator.run(
+                    step_id=1,
+                    display_step=1,
+                    lr=1e-3,
+                    save_checkpoint=save_ok,
+                )
+            kept = sorted(Path(tmpdir).glob("best.ckpt-*.pt"))
+            self.assertEqual(len(kept), 1)
+            self.assertIn("best.ckpt-1", kept[0].name)
+
+            with patch.object(
+                validator,
+                "evaluate_all_systems",
+                return_value={validator.metric_key: 0.1},
+            ):
+                with self.assertRaisesRegex(RuntimeError, "post-write"):
+                    validator.run(
+                        step_id=2,
+                        display_step=2,
+                        lr=1e-3,
+                        save_checkpoint=save_write_then_fail,
+                    )
+
+            self.assertEqual(len(validator.topk_records), 1)
+            self.assertEqual(validator.topk_records[0].step, 1)
+            remaining = sorted(Path(tmpdir).glob("best.ckpt-*.pt"))
+            self.assertEqual(len(remaining), 1)
+            self.assertIn("best.ckpt-1", remaining[0].name)
+            self.assertFalse(any("best.ckpt-2" in p.name for p in remaining))
+
     def test_distributed_remote_save_failure_rolls_back_committed_topk(self) -> None:
         """Rank 0 must not keep top-K when another rank aborts the save."""
         with tempfile.TemporaryDirectory() as tmpdir:

From d0ba5bc43e748df730790ffcec2ee0b4ce8c86c6 Mon Sep 17 00:00:00 2001
From: OutisLi-Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Sat, 10 Oct 2026 00:41:07 +0800
Subject: [PATCH 07/10] fix(train): keep best-checkpoint commit in distributed
 error phase

Chief-only commit/reconcile after a successful save could raise OSError and
exit rank 0 while other ranks waited in a later collective. Route commit
failures through _raise_if_distributed_error, and isolate cleanup errors so
they cannot bypass that collective.
---
 deepmd/pt_expt/train/validation.py | 36 +++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index 9441b95253..72aa86ecd2 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -367,8 +367,13 @@ def run(
                     # Match remote-failure cleanup: restore top-K, then drop any
                     # bytes already written for the aborted candidate (write-
                     # then-fail leaves an orphan that rollback alone keeps).
-                    self._rollback_pending_best_state()
-                    self._reconcile_best_checkpoints()
+                    # Swallow cleanup errors so they cannot bypass the
+                    # post-save collective below.
+                    try:
+                        self._rollback_pending_best_state()
+                        self._reconcile_best_checkpoints()
+                    except Exception:
+                        pass
                 caught_exception = exc
                 error_message = (
                     "Full validation failed while saving the best checkpoint:\n"
@@ -384,12 +389,31 @@ def run(
                 # Local save succeeded but another rank failed: undo staged
                 # top-K and drop the orphaned file before propagating.
                 if self.rank == 0 and caught_exception is None:
-                    self._rollback_pending_best_state()
-                    self._reconcile_best_checkpoints()
+                    try:
+                        self._rollback_pending_best_state()
+                        self._reconcile_best_checkpoints()
+                    except Exception:
+                        # Already propagating a remote failure; cleanup must
+                        # not replace that exception or skip non-chief waits.
+                        pass
                 raise
+            # Commit/reconcile stay inside a coordinated error phase so a
+            # chief-only OSError (e.g. prune/rename) cannot leave other ranks
+            # blocked in a later collective.
+            commit_error: str | None = None
+            commit_exception: Exception | None = None
             if self.rank == 0:
-                self._commit_pending_best_state()
-                self._reconcile_best_checkpoints()
+                try:
+                    self._commit_pending_best_state()
+                    self._reconcile_best_checkpoints()
+                except Exception as exc:
+                    commit_exception = exc
+                    commit_error = (
+                        "Full validation failed while committing the best "
+                        "checkpoint:\n"
+                        f"{traceback.format_exc()}"
+                    )
+            self._raise_if_distributed_error(commit_error, commit_exception)
 
         if self.rank == 0:
             try:

From d4ec233400b01b57ef831be0ac6d8ff3b2500712 Mon Sep 17 00:00:00 2001
From: OutisLi-Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Sat, 10 Oct 2026 00:48:11 +0800
Subject: [PATCH 08/10] fix(train): keep best-ckpt rollback until reconcile
 succeeds

Clear the rollback snapshot only after ranked rename finishes, restore
bookkeeping if reconcile raises, and undo partial .tmp renames so restart
can still see the checkpoint files.
---
 deepmd/pt_expt/train/validation.py | 75 ++++++++++++++++++------------
 1 file changed, 46 insertions(+), 29 deletions(-)

diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index 72aa86ecd2..347651e05b 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -406,7 +406,12 @@ def run(
                 try:
                     self._commit_pending_best_state()
                     self._reconcile_best_checkpoints()
+                    self._rollback_topk_records = None
                 except Exception as exc:
+                    try:
+                        self._rollback_pending_best_state()
+                    except Exception:
+                        pass
                     commit_exception = exc
                     commit_error = (
                         "Full validation failed while committing the best "
@@ -808,9 +813,12 @@ def _stage_pending_best_state(self) -> None:
         self._sync_state_store()
 
     def _commit_pending_best_state(self) -> None:
-        """Keep staged top-K after a successful checkpoint write."""
+        """Drop the proposal pointer after a successful checkpoint write.
+
+        The rollback snapshot stays until :meth:`_reconcile_best_checkpoints`
+        finishes, so a mid-reconcile failure can still restore bookkeeping.
+        """
         self._pending_topk_records = None
-        self._rollback_topk_records = None
 
     def _rollback_pending_best_state(self) -> None:
         """Restore pre-proposal top-K after a failed publication boundary."""
@@ -922,33 +930,42 @@ def _reconcile_best_checkpoints(self) -> None:
             files_by_step.setdefault(step, []).append(checkpoint_path)
 
         temp_moves: list[tuple[Path, Path]] = []
-        for step, checkpoint_paths in files_by_step.items():
-            expected_name = expected_names.get(step)
-            if expected_name is None:
-                stale_files.extend(checkpoint_paths)
-                continue
-
-            keep_path = next(
-                (
-                    checkpoint_path
-                    for checkpoint_path in checkpoint_paths
-                    if checkpoint_path.name == expected_name
-                ),
-                checkpoint_paths[0],
-            )
-            for checkpoint_path in checkpoint_paths:
-                if checkpoint_path != keep_path:
-                    stale_files.append(checkpoint_path)
-            if keep_path.name != expected_name:
-                temp_path = keep_path.with_name(f"{keep_path.name}.tmp")
-                keep_path.rename(temp_path)
-                temp_moves.append((temp_path, keep_path.with_name(expected_name)))
-
-        for checkpoint_path in stale_files:
-            self._remove_checkpoint_path(checkpoint_path)
-        for temp_path, final_path in temp_moves:
-            self._remove_checkpoint_path(final_path)
-            temp_path.rename(final_path)
+        try:
+            for step, checkpoint_paths in files_by_step.items():
+                expected_name = expected_names.get(step)
+                if expected_name is None:
+                    stale_files.extend(checkpoint_paths)
+                    continue
+
+                keep_path = next(
+                    (
+                        checkpoint_path
+                        for checkpoint_path in checkpoint_paths
+                        if checkpoint_path.name == expected_name
+                    ),
+                    checkpoint_paths[0],
+                )
+                for checkpoint_path in checkpoint_paths:
+                    if checkpoint_path != keep_path:
+                        stale_files.append(checkpoint_path)
+                if keep_path.name != expected_name:
+                    temp_path = keep_path.with_name(f"{keep_path.name}.tmp")
+                    keep_path.rename(temp_path)
+                    temp_moves.append((temp_path, keep_path.with_name(expected_name)))
+
+            for checkpoint_path in stale_files:
+                self._remove_checkpoint_path(checkpoint_path)
+            for temp_path, final_path in temp_moves:
+                self._remove_checkpoint_path(final_path)
+                temp_path.rename(final_path)
+        except Exception:
+            # Undo partial temp renames so restart can still see the files.
+            for temp_path, _final_path in temp_moves:
+                if temp_path.exists():
+                    original = temp_path.with_name(temp_path.name.removesuffix(".tmp"))
+                    if not original.exists():
+                        temp_path.rename(original)
+            raise
 
     def _initialize_best_checkpoints(self, restart_training: bool) -> None:
         """Align on-disk best checkpoints with the current training mode."""

From dd69afc85ea9dbd73cee5257b32598a43a138888 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Sat, 10 Oct 2026 00:57:39 +0800
Subject: [PATCH 09/10] fix(train): prune best-ckpt stales only after ranked
 renames
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_reconcile_best_checkpoints deleted stale files before finishing
temp→final renames. A mid-reconcile OSError after a stale delete left
_rollback_pending_best_state restoring the previous top-K while files
that top-K still named were already gone.

Complete ranked renames first, clear the rollback snapshot once the
on-disk retained set matches topk_records, then delete stales. Rename-
phase failures still undo .tmp moves and keep bookkeeping rollback.
Mirror the rename-then-delete order in dpmodel. Add regressions for
order, prune-fail, and rename-fail paths.
---
 deepmd/dpmodel/train/validation.py |  71 ++++---
 deepmd/pt_expt/train/validation.py |  30 ++-
 source/tests/pt/test_validation.py | 313 +++++++++++++++++++++++++++++
 3 files changed, 381 insertions(+), 33 deletions(-)

diff --git a/deepmd/dpmodel/train/validation.py b/deepmd/dpmodel/train/validation.py
index f652616be2..daf8a436fd 100644
--- a/deepmd/dpmodel/train/validation.py
+++ b/deepmd/dpmodel/train/validation.py
@@ -440,7 +440,12 @@ def _expected_topk_checkpoint_names(self) -> dict[int, str]:
         }
 
     def _reconcile_best_checkpoints(self) -> None:
-        """Rename retained best checkpoints to ranked names and delete stale ones."""
+        """Rename retained best checkpoints to ranked names and delete stale ones.
+
+        Ranked renames complete before any stale deletion so a mid-reconcile
+        failure cannot leave retained steps missing from disk while bookkeeping
+        still names them.
+        """
         expected_names = self._expected_topk_checkpoint_names()
         current_files = self._list_best_checkpoints()
         files_by_step: dict[int, list[Path]] = {}
@@ -454,33 +459,43 @@ def _reconcile_best_checkpoints(self) -> None:
             files_by_step.setdefault(step, []).append(checkpoint_path)
 
         temp_moves: list[tuple[Path, Path]] = []
-        for step, checkpoint_paths in files_by_step.items():
-            expected_name = expected_names.get(step)
-            if expected_name is None:
-                stale_files.extend(checkpoint_paths)
-                continue
-
-            keep_path = next(
-                (
-                    checkpoint_path
-                    for checkpoint_path in checkpoint_paths
-                    if checkpoint_path.name == expected_name
-                ),
-                checkpoint_paths[0],
-            )
-            for checkpoint_path in checkpoint_paths:
-                if checkpoint_path != keep_path:
-                    stale_files.append(checkpoint_path)
-            if keep_path.name != expected_name:
-                temp_path = keep_path.with_name(f"{keep_path.name}.tmp")
-                keep_path.rename(temp_path)
-                temp_moves.append((temp_path, keep_path.with_name(expected_name)))
-
-        for checkpoint_path in stale_files:
-            self._remove_checkpoint_path(checkpoint_path)
-        for temp_path, final_path in temp_moves:
-            self._remove_checkpoint_path(final_path)
-            temp_path.rename(final_path)
+        try:
+            for step, checkpoint_paths in files_by_step.items():
+                expected_name = expected_names.get(step)
+                if expected_name is None:
+                    stale_files.extend(checkpoint_paths)
+                    continue
+
+                keep_path = next(
+                    (
+                        checkpoint_path
+                        for checkpoint_path in checkpoint_paths
+                        if checkpoint_path.name == expected_name
+                    ),
+                    checkpoint_paths[0],
+                )
+                for checkpoint_path in checkpoint_paths:
+                    if checkpoint_path != keep_path:
+                        stale_files.append(checkpoint_path)
+                if keep_path.name != expected_name:
+                    temp_path = keep_path.with_name(f"{keep_path.name}.tmp")
+                    keep_path.rename(temp_path)
+                    temp_moves.append((temp_path, keep_path.with_name(expected_name)))
+
+            for temp_path, final_path in temp_moves:
+                self._remove_checkpoint_path(final_path)
+                temp_path.rename(final_path)
+            temp_moves.clear()
+
+            for checkpoint_path in stale_files:
+                self._remove_checkpoint_path(checkpoint_path)
+        except Exception:
+            for temp_path, _final_path in temp_moves:
+                if temp_path.exists():
+                    original = temp_path.with_name(temp_path.name.removesuffix(".tmp"))
+                    if not original.exists():
+                        temp_path.rename(original)
+            raise
 
     def _initialize_best_checkpoints(self, restart_training: bool) -> None:
         """Align on-disk best checkpoints with the current training mode."""
diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index 347651e05b..f4ec21c9f0 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -815,8 +815,10 @@ def _stage_pending_best_state(self) -> None:
     def _commit_pending_best_state(self) -> None:
         """Drop the proposal pointer after a successful checkpoint write.
 
-        The rollback snapshot stays until :meth:`_reconcile_best_checkpoints`
-        finishes, so a mid-reconcile failure can still restore bookkeeping.
+        The rollback snapshot stays until ranked renames in
+        :meth:`_reconcile_best_checkpoints` succeed, so a rename-phase failure
+        can still restore bookkeeping. It is cleared before stale deletion so a
+        mid-prune failure cannot restore records that name deleted files.
         """
         self._pending_topk_records = None
 
@@ -916,7 +918,15 @@ def _expected_topk_checkpoint_names(self) -> dict[int, str]:
         }
 
     def _reconcile_best_checkpoints(self) -> None:
-        """Rename retained best checkpoints to ranked names and delete stale ones."""
+        """Rename retained best checkpoints to ranked names and delete stale ones.
+
+        Ranked renames complete before any stale deletion. A failure during the
+        rename phase undoes ``.tmp`` moves so restart still sees the files, and
+        the caller may restore bookkeeping from the rollback snapshot. After
+        renames succeed the on-disk retained set matches ``topk_records``, so the
+        rollback snapshot is dropped before prune; a mid-delete failure must not
+        restore bookkeeping that names removed files.
+        """
         expected_names = self._expected_topk_checkpoint_names()
         current_files = self._list_best_checkpoints()
         files_by_step: dict[int, list[Path]] = {}
@@ -953,11 +963,21 @@ def _reconcile_best_checkpoints(self) -> None:
                     keep_path.rename(temp_path)
                     temp_moves.append((temp_path, keep_path.with_name(expected_name)))
 
-            for checkpoint_path in stale_files:
-                self._remove_checkpoint_path(checkpoint_path)
+            # Finalize ranked names before deleting anything still named by the
+            # previous top-K. Collision clears here only remove paths we are
+            # about to replace or that are already listed as stale.
             for temp_path, final_path in temp_moves:
                 self._remove_checkpoint_path(final_path)
                 temp_path.rename(final_path)
+            temp_moves.clear()
+
+            # Disk retained set now matches topk_records. Drop the snapshot so
+            # a mid-prune OSError cannot roll bookkeeping back to naming deleted
+            # stales; _rollback_pending_best_state becomes a no-op.
+            self._rollback_topk_records = None
+
+            for checkpoint_path in stale_files:
+                self._remove_checkpoint_path(checkpoint_path)
         except Exception:
             # Undo partial temp renames so restart can still see the files.
             for temp_path, _final_path in temp_moves:
diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py
index cd404a0cf8..1be39899fd 100644
--- a/source/tests/pt/test_validation.py
+++ b/source/tests/pt/test_validation.py
@@ -33,6 +33,7 @@
 from deepmd.pt_expt.train.validation import (
     BEST_METRIC_NAME_INFO_KEY,
     TOPK_RECORDS_INFO_KEY,
+    BestCheckpointRecord,
     FullValidator,
     resolve_full_validation_start_step,
 )
@@ -1373,3 +1374,315 @@ def broadcast_save_path(holder: list, src: int = 0) -> None:
             self.assertIsNone(result)
             self.assertEqual(len(writes), 1)
             self.assertTrue(writes[0].exists())
+
+    def test_reconcile_renames_before_deleting_stales(self) -> None:
+        """Ranked renames must finish before any stale prune."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = FullValidator(
+                validating_params={
+                    "full_validation": True,
+                    "validation_freq": 1,
+                    "save_best": True,
+                    "max_best_ckpt": 2,
+                    "validation_metric": "E:MAE",
+                    "full_val_file": str(Path(tmpdir) / "val.log"),
+                    "full_val_start": 0.0,
+                },
+                validation_data=_DummyValidationData(),
+                model=_DummyModel(),
+                state_store={},
+                num_steps=10,
+                rank=0,
+                restart_training=False,
+                checkpoint_dir=Path(tmpdir),
+            )
+            # Previous top-K on disk: steps 10 (rank1) and 20 (rank2).
+            Path(tmpdir, "best.ckpt-10.t-1.pt").write_text("10")
+            Path(tmpdir, "best.ckpt-20.t-2.pt").write_text("20")
+            # New top-K displaces 10; newly written 30 already at proposed rank.
+            Path(tmpdir, "best.ckpt-30.t-2.pt").write_text("30")
+            validator.topk_records = [
+                BestCheckpointRecord(metric=1.0, step=20),
+                BestCheckpointRecord(metric=1.5, step=30),
+            ]
+            validator._sync_state_store()
+
+            events: list[str] = []
+            real_rename = Path.rename
+            real_remove = FullValidator._remove_checkpoint_path
+
+            def tracking_rename(self_path: Path, target: Path) -> None:
+                events.append(f"rename:{self_path.name}->{target.name}")
+                real_rename(self_path, target)
+
+            def tracking_remove(path: Path) -> None:
+                events.append(f"remove:{path.name}")
+                real_remove(path)
+
+            with (
+                patch.object(Path, "rename", tracking_rename),
+                patch.object(
+                    FullValidator,
+                    "_remove_checkpoint_path",
+                    staticmethod(tracking_remove),
+                ),
+            ):
+                validator._reconcile_best_checkpoints()
+
+            remove_stale_idxs = [
+                i for i, e in enumerate(events) if e.startswith("remove:best.ckpt-10")
+            ]
+            final_rename_idxs = [
+                i
+                for i, e in enumerate(events)
+                if e.startswith("rename:") and e.endswith(".tmp->best.ckpt-20.t-1.pt")
+            ]
+            self.assertTrue(final_rename_idxs, events)
+            self.assertTrue(remove_stale_idxs, events)
+            self.assertLess(max(final_rename_idxs), min(remove_stale_idxs), events)
+            self.assertEqual(
+                sorted(path.name for path in Path(tmpdir).glob("best.ckpt-*.pt")),
+                ["best.ckpt-20.t-1.pt", "best.ckpt-30.t-2.pt"],
+            )
+
+    def test_prune_failure_does_not_restore_bookkeeping_naming_deleted_stales(
+        self,
+    ) -> None:
+        """After ranked renames, prune failure must not roll top-K back to deleted files.
+
+        Reproduces the pre-fix hazard: deleting ``best.ckpt-10…`` then failing
+        left rollback restoring steps ``{10,20}`` while disk only had ``{20,30}``.
+        """
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = FullValidator(
+                validating_params={
+                    "full_validation": True,
+                    "validation_freq": 1,
+                    "save_best": True,
+                    "max_best_ckpt": 2,
+                    "validation_metric": "E:MAE",
+                    "full_val_file": str(Path(tmpdir) / "val.log"),
+                    "full_val_start": 0.0,
+                },
+                validation_data=_DummyValidationData(),
+                model=_DummyModel(),
+                state_store={},
+                num_steps=10,
+                rank=0,
+                restart_training=False,
+                checkpoint_dir=Path(tmpdir),
+            )
+            old_records = [
+                BestCheckpointRecord(metric=1.0, step=10),
+                BestCheckpointRecord(metric=2.0, step=20),
+            ]
+            new_records = [
+                BestCheckpointRecord(metric=1.0, step=20),
+                BestCheckpointRecord(metric=1.5, step=30),
+            ]
+            Path(tmpdir, "best.ckpt-10.t-1.pt").write_text("10")
+            Path(tmpdir, "best.ckpt-20.t-2.pt").write_text("20")
+            Path(tmpdir, "best.ckpt-30.t-2.pt").write_text("30")
+            validator._rollback_topk_records = list(old_records)
+            validator.topk_records = list(new_records)
+            validator._sync_state_store()
+
+            real_remove = FullValidator._remove_checkpoint_path
+
+            def fail_on_stale_step10(path: Path) -> None:
+                if path.name.startswith("best.ckpt-10."):
+                    # Delete first (as a real OS would), then fail — matching the
+                    # mid-prune window the old ordering exposed to rollback.
+                    real_remove(path)
+                    raise OSError("simulated prune failure")
+                real_remove(path)
+
+            with patch.object(
+                FullValidator,
+                "_remove_checkpoint_path",
+                staticmethod(fail_on_stale_step10),
+            ):
+                with self.assertRaises(OSError):
+                    validator._reconcile_best_checkpoints()
+
+            # Snapshot cleared after renames; caller rollback is a no-op.
+            validator._rollback_pending_best_state()
+            named_steps = {record.step for record in validator.topk_records}
+            on_disk_steps = set()
+            for path in Path(tmpdir).glob("best.ckpt-*.pt"):
+                match = validator.best_checkpoint_pattern.match(path.name)
+                self.assertIsNotNone(match, path.name)
+                on_disk_steps.add(int(match.group(1)))
+            self.assertEqual(named_steps, {20, 30})
+            self.assertTrue(
+                named_steps.issubset(on_disk_steps),
+                f"bookkeeping {named_steps} missing on disk {on_disk_steps}",
+            )
+            self.assertEqual(
+                sorted(path.name for path in Path(tmpdir).glob("best.ckpt-*.pt")),
+                ["best.ckpt-20.t-1.pt", "best.ckpt-30.t-2.pt"],
+            )
+            self.assertIsNone(validator._rollback_topk_records)
+
+    def test_rename_phase_failure_keeps_rollback_and_restores_temps(self) -> None:
+        """Failure before ranked renames finish still allows bookkeeping rollback."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = FullValidator(
+                validating_params={
+                    "full_validation": True,
+                    "validation_freq": 1,
+                    "save_best": True,
+                    "max_best_ckpt": 2,
+                    "validation_metric": "E:MAE",
+                    "full_val_file": str(Path(tmpdir) / "val.log"),
+                    "full_val_start": 0.0,
+                },
+                validation_data=_DummyValidationData(),
+                model=_DummyModel(),
+                state_store={},
+                num_steps=10,
+                rank=0,
+                restart_training=False,
+                checkpoint_dir=Path(tmpdir),
+            )
+            old_records = [
+                BestCheckpointRecord(metric=1.0, step=10),
+                BestCheckpointRecord(metric=2.0, step=20),
+            ]
+            new_records = [
+                BestCheckpointRecord(metric=1.0, step=20),
+                BestCheckpointRecord(metric=1.5, step=30),
+            ]
+            Path(tmpdir, "best.ckpt-10.t-1.pt").write_text("10")
+            Path(tmpdir, "best.ckpt-20.t-2.pt").write_text("20")
+            Path(tmpdir, "best.ckpt-30.t-2.pt").write_text("30")
+            validator._rollback_topk_records = list(old_records)
+            validator.topk_records = list(new_records)
+            validator._sync_state_store()
+
+            real_rename = Path.rename
+
+            def fail_final_rename(self_path: Path, target: Path) -> None:
+                # Fail only the ranked finalization, not the .tmp→original undo.
+                if self_path.name.endswith(".tmp"):
+                    original_name = self_path.name.removesuffix(".tmp")
+                    if target.name != original_name:
+                        raise OSError("simulated ranked rename failure")
+                real_rename(self_path, target)
+
+            with patch.object(Path, "rename", fail_final_rename):
+                with self.assertRaises(OSError):
+                    validator._reconcile_best_checkpoints()
+
+            # Temps undone; stale 10 still present; rollback snapshot retained.
+            self.assertIsNotNone(validator._rollback_topk_records)
+            validator._rollback_pending_best_state()
+            named_steps = {record.step for record in validator.topk_records}
+            on_disk_steps = set()
+            for path in Path(tmpdir).glob("best.ckpt-*.pt"):
+                match = validator.best_checkpoint_pattern.match(path.name)
+                self.assertIsNotNone(match, path.name)
+                on_disk_steps.add(int(match.group(1)))
+            self.assertEqual(named_steps, {10, 20})
+            self.assertTrue(
+                named_steps.issubset(on_disk_steps),
+                f"bookkeeping {named_steps} missing on disk {on_disk_steps}",
+            )
+            self.assertFalse(list(Path(tmpdir).glob("*.tmp")))
+
+    def test_commit_path_prune_failure_keeps_named_checkpoints(self) -> None:
+        """run() commit/reconcile prune failure must not orphan named top-K steps."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            validator = FullValidator(
+                validating_params={
+                    "full_validation": True,
+                    "validation_freq": 1,
+                    "save_best": True,
+                    "max_best_ckpt": 2,
+                    "validation_metric": "E:MAE",
+                    "full_val_file": str(Path(tmpdir) / "val.log"),
+                    "full_val_start": 0.0,
+                },
+                validation_data=_DummyValidationData(),
+                model=_DummyModel(),
+                state_store={},
+                num_steps=10,
+                rank=0,
+                restart_training=False,
+                checkpoint_dir=Path(tmpdir),
+            )
+
+            def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None:
+                del lr, step
+                Path(path).write_text(f"ckpt-{path.name}")
+
+            # Seed previous best pair via a successful run, then force a third
+            # candidate so reconcile must prune the displaced step.
+            with patch.object(
+                validator,
+                "evaluate_all_systems",
+                return_value={validator.metric_key: 2.0},
+            ):
+                validator.run(
+                    step_id=10,
+                    display_step=10,
+                    lr=1e-3,
+                    save_checkpoint=save_checkpoint,
+                )
+            with patch.object(
+                validator,
+                "evaluate_all_systems",
+                return_value={validator.metric_key: 1.0},
+            ):
+                validator.run(
+                    step_id=20,
+                    display_step=20,
+                    lr=1e-3,
+                    save_checkpoint=save_checkpoint,
+                )
+            self.assertEqual(
+                {record.step for record in validator.topk_records},
+                {10, 20},
+            )
+
+            real_remove = FullValidator._remove_checkpoint_path
+
+            def fail_on_stale_step10(path: Path) -> None:
+                if path.name.startswith("best.ckpt-10."):
+                    real_remove(path)
+                    raise OSError("simulated prune failure")
+                real_remove(path)
+
+            with (
+                patch.object(
+                    validator,
+                    "evaluate_all_systems",
+                    return_value={validator.metric_key: 1.5},
+                ),
+                patch.object(
+                    FullValidator,
+                    "_remove_checkpoint_path",
+                    staticmethod(fail_on_stale_step10),
+                ),
+            ):
+                with self.assertRaisesRegex(RuntimeError, "committing the best"):
+                    validator.run(
+                        step_id=30,
+                        display_step=30,
+                        lr=1e-3,
+                        save_checkpoint=save_checkpoint,
+                    )
+
+            named_steps = {record.step for record in validator.topk_records}
+            on_disk_steps = set()
+            for path in Path(tmpdir).glob("best.ckpt-*.pt"):
+                match = validator.best_checkpoint_pattern.match(path.name)
+                self.assertIsNotNone(match, path.name)
+                on_disk_steps.add(int(match.group(1)))
+            self.assertEqual(named_steps, {20, 30})
+            self.assertTrue(
+                named_steps.issubset(on_disk_steps),
+                f"bookkeeping {named_steps} missing on disk {on_disk_steps}",
+            )
+            self.assertIsNone(validator._rollback_topk_records)
+            self.assertIsNone(validator._pending_topk_records)

From d25cd7d582950cc543cfeb23ed1ae420b226abd2 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Sat, 10 Oct 2026 10:42:45 +0800
Subject: [PATCH 10/10] fix(train): document intentional empty excepts in
 best-ckpt cleanup

CodeQL EmptyExcept on the local-save and commit-failure cleanup paths;
swallowing nested rollback errors is required so the primary exception
still reaches the post-save collective / commit error phase.
---
 deepmd/pt_expt/train/validation.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index f4ec21c9f0..2362e532be 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -367,12 +367,13 @@ def run(
                     # Match remote-failure cleanup: restore top-K, then drop any
                     # bytes already written for the aborted candidate (write-
                     # then-fail leaves an orphan that rollback alone keeps).
-                    # Swallow cleanup errors so they cannot bypass the
-                    # post-save collective below.
                     try:
                         self._rollback_pending_best_state()
                         self._reconcile_best_checkpoints()
                     except Exception:
+                        # Swallow cleanup errors so they cannot bypass the
+                        # post-save collective below (primary save failure
+                        # must still be broadcast to every rank).
                         pass
                 caught_exception = exc
                 error_message = (
@@ -411,6 +412,8 @@ def run(
                     try:
                         self._rollback_pending_best_state()
                     except Exception:
+                        # Already recording a commit failure; nested rollback
+                        # errors must not replace that exception.
                         pass
                     commit_exception = exc
                     commit_error = (