Repository: allenai/tango (language: python)

Issue:
Fix Trainer crash



Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 187178840..132501ed4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -40,6 +40,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Fixed bug where `StepInfo` environment and platform metadata could be out-of-date if a step is run again due to failure.
+- Fixed a bug where an unfortunate combination of early stopping and decreasing model performance could result in a crash in the torch trainer.
 
 ## [v0.11.0](https://github.com/allenai/tango/releases/tag/v0.11.0) - 2022-08-04
 
diff --git a/tango/integrations/torch/train.py b/tango/integrations/torch/train.py
index 014b09bac..edd0d8714 100644
--- a/tango/integrations/torch/train.py
+++ b/tango/integrations/torch/train.py
@@ -268,6 +268,12 @@ def _train(
                 "One of 'train_steps' or 'train_epochs' needs to be specified, but not both."
             )
 
+        if validate_every is not None and checkpoint_every is not None:
+            if checkpoint_every % validate_every != 0 and validate_every % checkpoint_every != 0:
+                raise ConfigurationError(
+                    "'checkpoint_every' needs to be multiple of 'validate_every' or vice versa"
+                )
+
         config = TrainConfig(
             self.unique_id,
             self.work_dir,
@@ -424,15 +430,27 @@ def _train(
         if validation_dataloader is not None:
             check_dataloader(validation_dataloader)
 
+    # The (training) loss for each batch, updated every training batch.
     batch_loss: float = 0.0
-    best_batch_loss: Optional[float] = None
+    # The value of the validation metric (could be loss), updated after every validation loop.
     val_metric: Optional[float] = None
+    # The best validation metric over all validation set passes.
     best_val_metric: Optional[float] = None
+    # The best validation metric over all validation set passes that correspond to a checkpoint.
+    # Could be different from `best_val_metric` if `checkpoint_every` > `validate_every`.
+    best_val_metric_checkpointed: Optional[float] = None
+    # The step to start training from.
     start_step: int = 0
+    # The current training step.
+    step: int = start_step
+    # If we should do a validation pass after the current training batch.
+    should_validate_this_step: bool = False
+
+    # Load state from checkpoint.
     if initial_state is not None:
         val_metric = initial_state[f"val_{config.val_metric_name}"]
         best_val_metric = initial_state[f"best_{config.val_metric_name}"]
-        best_batch_loss = initial_state["best_batch_loss"]
+        best_val_metric_checkpointed = initial_state[f"best_{config.val_metric_name}_checkpointed"]
         start_step = initial_state["training_steps"]
 
     # Initialize callbacks.
@@ -463,32 +481,45 @@ def _train(
 
     def is_best_checkpoint() -> bool:
         """
-        A closure that we'll call when saving checkpoints to check if we should hardlink
+        A closure that we'll call when saving checkpoints to check if we should link
         the best checkpoint path to the current checkpoint file.
         """
-        if val_metric is not None and best_val_metric is not None:
-            return (config.minimize_val_metric and val_metric <= best_val_metric) or (
-                not config.minimize_val_metric and val_metric >= best_val_metric
-            )
-        elif best_batch_loss is not None:
-            return best_batch_loss <= batch_loss
+        if val_metric is not None:
+            if best_val_metric_checkpointed is not None:
+                return (
+                    config.minimize_val_metric and val_metric <= best_val_metric_checkpointed
+                ) or (not config.minimize_val_metric and val_metric >= best_val_metric_checkpointed)
+            else:
+                return False
         else:
-            return False
+            # Without a validation loop we always treat the most recent checkpoint as the best.
+            return True
 
     def save_state(step: int):
         """
         A closure that we'll call every `checkpoint_every` steps in the train loop to
         save model and training state.
         """
+        # Update best loss/metric trackers.
+        nonlocal best_val_metric_checkpointed
+        if should_validate_this_step and val_metric is not None:
+            if (
+                best_val_metric_checkpointed is None
+                or (config.minimize_val_metric and val_metric <= best_val_metric_checkpointed)
+                or (not config.minimize_val_metric and val_metric >= best_val_metric_checkpointed)
+            ):
+                best_val_metric_checkpointed = val_metric
+
         train_state = {
             "training_steps": step + 1,
-            "best_batch_loss": best_batch_loss,
             f"val_{config.val_metric_name}": val_metric,
             f"best_{config.val_metric_name}": best_val_metric,
+            f"best_{config.val_metric_name}_checkpointed": best_val_metric_checkpointed,
             "callbacks": [
                 callback.state_dict() for callback in callbacks  # type: ignore[union-attr]
             ],
         }
+
         # For reason mypy can't figure out that `training_engine` is a `TrainingEngine` in this closure,
         # and not a `Lazy[TrainingEngine]`.
         cast(TrainingEngine, training_engine).save_checkpoint(
@@ -595,9 +626,6 @@ def save_state(step: int):
             for callback in callbacks:
                 callback.post_batch(step, current_epoch, batch_loss)
 
-            if best_batch_loss is None or batch_loss <= best_batch_loss:
-                best_batch_loss = batch_loss
-
             del batch
 
             training_engine.step()
@@ -605,22 +633,24 @@ def save_state(step: int):
             # Find out whether we should validate
             if config.validation_split is None:
                 # If we can't validate, we don't.
-                should_validate = False
+                should_validate_this_step = False
             elif step == config.train_steps - 1:
                 # If we're at the end of the training run, we always validate.
-                should_validate = True
+                should_validate_this_step = True
             elif config.validate_every is not None and (step + 1) % config.validate_every == 0:
                 # If validate_every is given, we use that to decide.
-                should_validate = True
+                should_validate_this_step = True
             elif config.validate_every is None and epoch != train_batch_iterator.peek()[1][0]:
                 # If validate_every is not given, we validate at the end of the epoch.
-                should_validate = True
+                should_validate_this_step = True
             else:
                 # Otherwise, we don't validate.
-                should_validate = False
+                should_validate_this_step = False
 
             # Gather average loss across all workers.
-            if (config.should_log_this_step(step) or should_validate) and config.is_distributed:
+            if (
+                config.should_log_this_step(step) or should_validate_this_step
+            ) and config.is_distributed:
                 batch_loss_tensor = torch.tensor(batch_loss, device=device)
                 dist.all_reduce(batch_loss_tensor)
                 batch_loss = batch_loss_tensor.item() / config.world_size
@@ -640,7 +670,7 @@ def save_state(step: int):
                     train_batch_iterator_tqdm.set_postfix(**metrics_to_log)
 
             # Validate.
-            if should_validate:
+            if should_validate_this_step:
                 assert validation_dataloader is not None
                 assert config.validation_steps is not None
 
@@ -698,17 +728,25 @@ def save_state(step: int):
                 # Reset model to train mode.
                 training_engine.model.train()
 
-                if best_val_metric is None:
-                    best_val_metric = val_metric
-                elif config.minimize_val_metric and val_metric <= best_val_metric:
-                    best_val_metric = val_metric
-                elif not config.minimize_val_metric and val_metric >= best_val_metric:
+                if (
+                    best_val_metric is None
+                    or (config.minimize_val_metric and val_metric <= best_val_metric)
+                    or (not config.minimize_val_metric and val_metric >= best_val_metric)
+                ):
                     best_val_metric = val_metric
 
+                # Checkpoint.
+                if config.should_checkpoint_this_step(step):
+                    save_state(step)
+
                 # Post validation callback.
                 for callback in callbacks:
                     callback.post_val_loop(step, current_epoch, val_metric, best_val_metric)
 
+                # Reset model to train mode again in case the callbacks messed with it.
+                if callbacks:
+                    training_engine.model.train()
+
                 # Update progress bar again.
                 metrics_to_log = {
                     "batch_loss": batch_loss,
@@ -717,10 +755,10 @@ def save_state(step: int):
                 }
                 if config.is_local_main_process:
                     train_batch_iterator_tqdm.set_postfix(**metrics_to_log)
-
-            # Checkpoint.
-            if config.should_checkpoint_this_step(step):
-                save_state(step)
+            else:
+                # Checkpoint.
+                if config.should_checkpoint_this_step(step):
+                    save_state(step)
 
         # End train loop
 
@@ -736,6 +774,10 @@ def save_state(step: int):
     if config.is_distributed:
         dist.barrier()
 
+    # If we haven't saved a checkpoint yet, do it now.
+    if not config.best_state_path.exists():
+        save_state(step)
+
     for callback in callbacks:
         callback.post_train_loop(step, current_epoch)
 


Test patch:
diff --git a/tests/integrations/torch/training_engine_test.py b/tests/integrations/torch/training_engine_test.py
index c6ff4ec67..fb2318c26 100644
--- a/tests/integrations/torch/training_engine_test.py
+++ b/tests/integrations/torch/training_engine_test.py
@@ -1,8 +1,18 @@
+import time
+
 import pytest
 import torch
 import torch.nn as nn
+from torch.nn import MSELoss
 
+from tango.common import DatasetDict, Lazy
 from tango.common.testing import TangoTestCase
+from tango.integrations.torch import (
+    DataLoader,
+    StopEarly,
+    TorchTrainStep,
+    TrainCallback,
+)
 from tango.integrations.torch.model import Model
 from tango.integrations.torch.training_engine import TorchTrainingEngine
 
@@ -39,3 +49,49 @@ def test_grad_scaler(self):
 
         assert (self.TEST_DIR / "worker0_grad_scaler.pt").is_file()
         assert training_engine.grad_scaler == saved_grad_scaler
+
+
+class WorseningModel(Model):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(7, 1)
+        self.loss = MSELoss()
+        self.start_time = time.time()
+
+    def forward(self, x, y):
+        y_hat = self.linear(x)
+        time.sleep(0.01)
+        return {"loss": self.loss(y_hat, y) + (time.time() - self.start_time)}
+
+
+class StopOnStepCallback(TrainCallback):
+    def __init__(self, stop_on_step: int, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.stop_on_step = stop_on_step
+
+    def post_val_loop(
+        self, step: int, epoch: int, val_metric: float, best_val_metric: float
+    ) -> None:
+        if step == self.stop_on_step:
+            raise StopEarly
+
+
+def test_with_increasing_loss():
+    model = WorseningModel()
+
+    xs = [torch.randn(7) for _ in range(100)]
+    train_set = [{"x": x, "y": x + 0.1} for x in xs]
+    dataset = DatasetDict(splits={"train": train_set, "validation": train_set}, metadata={})
+
+    step = TorchTrainStep(
+        model=model,
+        training_engine=Lazy(TorchTrainingEngine, optimizer=Lazy(torch.optim.AdamW, lr=1e-5)),
+        dataset_dict=dataset,
+        train_dataloader=Lazy(DataLoader),
+        train_steps=10,
+        validation_steps=10,
+        train_split="train",
+        validation_split="validation",
+        callbacks=[Lazy(StopOnStepCallback, stop_on_step=9)],
+    )
+    step.result()