fix(train): replay epochs after prevalidation rollback

tests/test_engine.py

@@ -324,11 +324,14 @@ def poison_ema(trainer):
 
 
 def test_checkpoint_nonfinite_ema_and_model_sanitized():
-    """Test a tensor non-finite in both EMA and model is sanitized (not skipped) so the run still produces a checkpoint."""
+    """A one-time online/EMA fault is restored and replayed before saving a completed epoch."""
+    injected = False
 
     def poison_ema_and_model(trainer):
         """Force the first parameter non-finite in both the live EMA and the model (finite-loss sticky-NaN)."""
-        if trainer.ema is not None:
+        nonlocal injected
+        if trainer.ema is not None and not injected:
+            injected = True
             next(iter(trainer.ema.ema.parameters())).data.flatten()[0] = float("inf")
             next(iter(unwrap_model(trainer.model).parameters())).data.flatten()[0] = float("nan")
 

tests/test_prevalidation_recovery.py

@@ -0,0 +1,92 @@
+"""Exercise epoch replay when callbacks corrupt online state before validation."""
+
+from collections import Counter
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from ultralytics.models.yolo.detect import DetectionTrainer
+from ultralytics.utils.torch_utils import unwrap_model
+
+
+@pytest.mark.parametrize("persistent", [False, True])
+def test_prevalidation_recovery_replays_or_exhausts_budget(tmp_path, persistent):
+    """Discard rolled-back attempts, preserve recovery limits and reset optimizer accumulation."""
+    for split in ("train", "val"):
+        images, labels = tmp_path / "images" / split, tmp_path / "labels" / split
+        images.mkdir(parents=True)
+        labels.mkdir(parents=True)
+        for index in range(4):
+            pixels = np.full((32, 32, 3), 40 + index * 30, dtype=np.uint8)
+            Image.fromarray(pixels).save(images / f"{index}.jpg")
+            (labels / f"{index}.txt").write_text("0 0.5 0.5 0.5 0.5\n")
+    data = tmp_path / "data.yaml"
+    data.write_text(f"path: {tmp_path.as_posix()}\ntrain: images/train\nval: images/val\nnames: [object]\n")
+    trainer = DetectionTrainer(
+        overrides={
+            "model": "yolo11n.yaml",
+            "data": str(data),
+            "imgsz": 32,
+            "epochs": 1,
+            "batch": 2,
+            "nbs": 2,
+            "workers": 0,
+            "device": "cpu",
+            "amp": False,
+            "mosaic": 0.0,
+            "close_mosaic": 0,
+            "optimizer": "SGD",
+            "warmup_epochs": 0,
+            "plots": False,
+            "project": str(tmp_path / "runs"),
+            "name": "recovery",
+        }
+    )
+    trainer.final_eval = lambda: None
+    attempts, finalized, validated, saved = [], [], [], []
+    steps = Counter()
+    original_step, original_validate = trainer.optimizer_step, trainer.validate
+    original_save, original_finalize = trainer.save_metrics, trainer._finalize_moe_map_saturation_epoch
+
+    def step():
+        steps[len(attempts)] += 1
+        return original_step()
+
+    def validate():
+        validated.append(len(attempts))
+        return original_validate()
+
+    def save(metrics):
+        saved.append((len(attempts), trainer.epoch))
+        return original_save(metrics)
+
+    def finalize(**kwargs):
+        finalized.append(kwargs)
+        return original_finalize(**kwargs)
+
+    def poison(t):
+        if persistent or len(attempts) == 1:
+            with torch.no_grad():
+                next(unwrap_model(t.model).parameters()).flatten()[0] = float("nan")
+
+    trainer.optimizer_step, trainer.validate = step, validate
+    trainer.save_metrics, trainer._finalize_moe_map_saturation_epoch = save, finalize
+    trainer.add_callback("on_train_epoch_start", lambda t: attempts.append(t.epoch))
+    trainer.add_callback("on_train_epoch_end", poison)
+    if persistent:
+        with pytest.raises(RuntimeError, match="NaN persisted"):
+            trainer.train()
+        assert attempts == [0, 0, 0, 0]
+        assert trainer.nan_recovery_attempts == 4
+        assert not saved and not validated
+        assert finalized == [{"recovered": True, "validated": False}] * 3
+    else:
+        trainer.train()
+        assert attempts == [0, 0]
+        assert steps == {1: 2, 2: 2}
+        assert validated == [2] and saved == [(2, 0)]
+        assert finalized == [{"recovered": True, "validated": False}, {"recovered": False, "validated": True}]
+        assert trainer.nan_recovery_attempts == 0
+        assert len(trainer.csv.read_text().splitlines()) == 2

ultralytics/engine/trainer.py

@@ -822,7 +822,8 @@ def _do_train(self):
             if validated:
                 self._clear_memory(None if self.device.type == "mps" else 0.5)  # prevent VRAM spike
                 if self._recover_before_validation(epoch):
-                    self._finalize_moe_map_saturation_epoch(recovered=True, validated=True)
+                    last_opt_step = _reset_optimizer_accumulation_after_recovery(self.optimizer)
+                    self._finalize_moe_map_saturation_epoch(recovered=True, validated=False)
                     continue
                 self.metrics, self.fitness = self.validate()
 
@@ -1220,7 +1221,7 @@ def _collect_prevalidation_nonfinite_flags(self):
         }
 
     def _recover_before_validation(self, epoch):
-        """Recover before validation if the online or EMA model is already non-finite."""
+        """Return whether to replay the epoch after prevalidation recovery; EMA-only resync needs no replay."""
         flags = self._collect_prevalidation_nonfinite_flags()
         if flags["ema_nonfinite"]:
             self._recovery_controller().resync_nonfinite_ema()
@@ -1230,9 +1231,9 @@ def _recover_before_validation(self, epoch):
         self.fitness = float("nan")
         recovered = self._handle_nan_recovery(epoch)
         if recovered:
-            # The live graph is finite again. Validate and checkpoint the restored state
-            # instead of replaying an epoch that may repeat a deterministic callback fault.
-            return False
+            # Rolling back the online model discards this attempt's updates. Replay it without
+            # committing stale metrics or clearing the consecutive recovery budget.
+            return True
         return any(self._collect_prevalidation_nonfinite_flags().values())
 
     def _record_nonfinite_diagnostic(self, component, *, epoch, step, loss_items=None, parameter=None):