fix(train): gate every checkpoint publish with NonFiniteGradGuard

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."""

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,

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(