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/train/training.py b/deepmd/pt/train/training.py index 7efc76e4d9..32f8b8b962 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -1732,22 +1732,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 ( @@ -1756,11 +1748,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 ): @@ -2032,6 +2021,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 f1e7f27c58..16babecf49 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -3058,10 +3058,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: @@ -3083,6 +3094,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( @@ -3097,6 +3109,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..2362e532be 100644 --- a/deepmd/pt_expt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -289,6 +289,8 @@ 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) @@ -344,14 +346,35 @@ 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: - # 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) + # 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 every + # rank reports a successful boundary. if self.rank == 0: - self._reconcile_best_checkpoints() + 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) 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). + 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 = ( "Full validation failed while saving the best checkpoint:\n" @@ -361,7 +384,44 @@ 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: + 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: + 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: + # Already recording a commit failure; nested rollback + # errors must not replace that exception. + pass + 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: @@ -408,12 +468,22 @@ 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 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]) - 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 +766,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 +792,68 @@ 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 _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() - 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)) + + def _commit_pending_best_state(self) -> None: + """Drop the proposal pointer after a successful checkpoint write. + + 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 + + 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, + *, + 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 be staged for the + write and rolled back if publication fails. + """ + 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.""" @@ -789,7 +921,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]] = {} @@ -803,33 +943,52 @@ 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))) + + # 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: + 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/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py index 767b9967e4..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, ) @@ -990,3 +991,698 @@ 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] = [] + 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)) + + 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) + 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: + 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.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_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: + 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: + 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()) + + 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) diff --git a/source/tests/pt_expt/test_train_gradient.py b/source/tests/pt_expt/test_train_gradient.py index f20912f57b..8434cdde6e 100644 --- a/source/tests/pt_expt/test_train_gradient.py +++ b/source/tests/pt_expt/test_train_gradient.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest +from pathlib import ( + Path, +) import torch @@ -111,5 +114,100 @@ def test_resets_after_check(self) -> None: guard.raise_if_nonfinite(self._named(1.0)) +class TestCheckpointPublicationGuard(unittest.TestCase): + """Publication-boundary contract via real Trainer save entry points (#5816).""" + + @staticmethod + 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 _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): + trainer._save_full_validation_checkpoint(Path("best.pt"), step=3) + self.assertEqual(writes, []) + + def test_ema_validation_best_raises_before_serialize(self) -> None: + trainer = self._bare_trainer() + trainer.nonfinite_grad_guard.update(torch.tensor(float("inf"), device="cpu")) + writes: list[object] = [] + + 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): + trainer._save_full_validation_ema_checkpoint(Path("best-ema.pt"), step=4) + self.assertEqual(writes, []) + + 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[bool] = [] + real_ensure = trainer.ensure_finite_gradients_for_checkpoint + + def counting_ensure() -> None: + checks["count"] += 1 + real_ensure() + + def fake_write(path, *, step: int, use_ema_weights: bool = False) -> None: + del path, step + writes.append(use_ema_weights) + + 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, [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) + + if __name__ == "__main__": unittest.main()