Allow disabling automatic stopping after max_steps or max_epochs

docs/source/common/trainer.rst

@@ -907,6 +907,9 @@ Stop training once this number of epochs is reached
     # default used by the Trainer
     trainer = Trainer(max_epochs=1000)
 
+If both ``max_epochs`` and ``max_steps`` aren't specified, ``max_epochs`` will default to ``1000``.
+To disable this default, set ``max_epochs = -1``.
+
 min_epochs
 ^^^^^^^^^^
 
@@ -947,6 +950,9 @@ Training will stop if max_steps or max_epochs have reached (earliest).
     # Stop after 100 steps
     trainer = Trainer(max_steps=100)
 
+If ``max_steps`` is not specified, ``max_epochs`` will be used instead (and ``max_epochs`` defaults to
+``1000`` if ``max_epochs`` is not specified). To disable this default, set ``max_steps = -1``.
+
 min_steps
 ^^^^^^^^^
 

pytorch_lightning/loops/fit_loop.py

@@ -124,6 +124,19 @@ def _results(self) -> ResultCollection:
             return self.epoch_loop.val_loop._results
         raise RuntimeError("`FitLoop._results` property isn't defined. Accessed outside of scope")
 
+    @staticmethod
+    def _is_max_limit_enabled(max_value: Optional[int]) -> bool:
+        """Checks whether the max_value is enabled. This can
+        be used for checking whether max_epochs or max_steps is enabled.
+
+        Args:
+            max_value: the value to check
+
+        Returns:
+            whether the limit for this value should be enabled
+        """
+        return max_value not in (None, -1)
+
     @property
     def done(self) -> bool:
         """Evaluates when to leave the loop.
@@ -132,8 +145,8 @@ def done(self) -> bool:
         or if the maximum number of steps or epochs is reached.
         """
         # TODO(@awaelchli): Move track steps inside training loop and move part of these condition inside training loop
-        stop_steps = self.max_steps is not None and self.global_step >= self.max_steps
-        stop_epochs = self.max_epochs is not None and self.current_epoch >= self.max_epochs
+        stop_steps = FitLoop._is_max_limit_enabled(self.max_steps) and self.global_step >= self.max_steps
+        stop_epochs = FitLoop._is_max_limit_enabled(self.max_epochs) and self.current_epoch >= self.max_epochs
 
         should_stop = False
         if self.trainer.should_stop:

pytorch_lightning/trainer/configuration_validator.py

@@ -90,6 +90,17 @@ def __verify_train_loop_configuration(self, model: "pl.LightningModule") -> None
                 "(rather, they are called on every optimization step)."
             )
 
+        # Allow max_epochs or max_steps to be zero, since this will be handled by fit_loop.done
+        if trainer.max_epochs and (trainer.max_epochs < -1 or not isinstance(trainer.max_epochs, int)):
+            raise MisconfigurationException(
+                f"`max_epochs` must be a positive integer or -1. You passed in {trainer.max_epochs}."
+            )
+
+        if trainer.max_steps and (trainer.max_steps < -1 or not isinstance(trainer.max_steps, int)):
+            raise MisconfigurationException(
+                f"`max_steps` must be a positive integer or -1. You passed in {trainer.max_steps}."
+            )
+
     def __verify_eval_loop_configuration(self, model: "pl.LightningModule", stage: str) -> None:
         loader_name = f"{stage}_dataloader"
         step_name = "validation_step" if stage == "val" else "test_step"

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -21,6 +21,7 @@
 from torchmetrics import Metric
 
 import pytorch_lightning as pl
+from pytorch_lightning.loops.fit_loop import FitLoop
 from pytorch_lightning.utilities import _OMEGACONF_AVAILABLE, rank_zero_deprecation, rank_zero_info, rank_zero_warn
 from pytorch_lightning.utilities.cloud_io import atomic_save, get_filesystem
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
@@ -190,7 +191,10 @@ def restore_loops(self) -> None:
         self.trainer.fit_loop.current_epoch = self._loaded_checkpoint["epoch"]
 
         # crash if max_epochs is lower then the current epoch from the checkpoint
-        if self.trainer.max_epochs is not None and self.trainer.current_epoch > self.trainer.max_epochs:
+        if (
+            FitLoop._is_max_limit_enabled(self.trainer.max_epochs)
+            and self.trainer.current_epoch > self.trainer.max_epochs
+        ):
             raise MisconfigurationException(
                 f"You restored a checkpoint with current_epoch={self.trainer.current_epoch},"
                 f" but you have set Trainer(max_epochs={self.trainer.max_epochs})."

pytorch_lightning/trainer/trainer.py

@@ -261,12 +261,15 @@ def __init__(
                 Can be used on CPU, GPU or TPUs.
 
             max_epochs: Stop training once this number of epochs is reached. Disabled by default (None).
-                If both max_epochs and max_steps are not specified, defaults to ``max_epochs`` = 1000.
+                If both max_epochs and max_steps are not specified, defaults to ``max_epochs = 1000``.
+                To disable this default, set ``max_epochs = -1``.
 
             min_epochs: Force training for at least these many epochs. Disabled by default (None).
-                If both min_epochs and min_steps are not specified, defaults to ``min_epochs`` = 1.
+                If both min_epochs and min_steps are not specified, defaults to ``min_epochs = 1``.
 
-            max_steps: Stop training after this number of steps. Disabled by default (None).
+            max_steps: Stop training after this number of steps. Disabled by default (None). If ``max_steps = None``
+                and ``max_epochs = None``, will default to ``max_epochs = 1000``. To disable this default, set
+                ``max_steps`` to ``-1``.
 
             min_steps: Force training for at least these number of steps. Disabled by default (None).
 
@@ -374,6 +377,7 @@ def __init__(
         self.slurm_connector = SLURMConnector(self)
         self.tuner = Tuner(self)
 
+        # max_epochs won't default to 1000 if max_steps/max_time are specified (including being set to -1).
         fit_loop = FitLoop(
             min_epochs=(1 if (min_epochs is None and min_steps is None and max_time is None) else min_epochs),
             max_epochs=(1000 if (max_epochs is None and max_steps is None and max_time is None) else max_epochs),

tests/callbacks/test_timer.py

@@ -42,8 +42,12 @@ def on_fit_start(self):
         trainer.fit(TestModel())
     assert "callbacks list already contains a Timer" in caplog.text
 
-    seconds = 1
-    trainer = Trainer(max_time=dict(seconds=seconds))
+    # Make sure max_time still honored even if max_epochs == -1
+    trainer = Trainer(max_time=dict(seconds=1), max_epochs=-1)
+    with pytest.raises(SystemExit):
+        trainer.fit(TestModel())
+    timer = [c for c in trainer.callbacks if isinstance(c, Timer)][0]
+    assert timer._duration == 1
     assert trainer.max_epochs is None
     assert trainer.max_steps is None
 

tests/trainer/test_trainer.py

@@ -491,6 +491,56 @@ def test_trainer_max_steps_and_epochs(tmpdir):
     assert trainer.global_step == num_train_samples * trainer.max_epochs
     assert trainer.current_epoch == trainer.max_epochs - 1, "Model did not stop at max_epochs"
 
+    # if max_steps is positive and max_epochs is negative, use max_steps
+    trainer_kwargs["max_epochs"] = -1
+    trainer_kwargs["max_steps"] = 3 * 2 * num_train_samples
+    trainer = Trainer(**trainer_kwargs)
+    trainer.fit(model)
+
+    assert trainer.state.finished, f"Training failed with {trainer.state}"
+    assert trainer.global_step == 3 * 2 * num_train_samples
+
+
+@pytest.mark.parametrize(
+    "max_epochs,max_steps,incorrect_variable,incorrect_value",
+    [
+        (-100, None, "max_epochs", -100),
+        (1.5, None, "max_epochs", 1.5),
+        (1, -2, "max_steps", -2),
+        (1, 0.5, "max_steps", -2),
+    ],
+)
+def test_trainer_max_steps_and_epochs_validation(max_epochs, max_steps, incorrect_variable, incorrect_value):
+    """Don't allow max_epochs or max_steps to be less than -1 or a float"""
+    with pytest.raises(
+        MisconfigurationException,
+        match=f"`{incorrect_variable}` must be a positive integer or -1. You passed in {incorrect_value}",
+    ):
+        trainer = Trainer(max_epochs=max_epochs, max_steps=max_steps)
+
+
+@pytest.mark.parametrize(
+    "max_epochs,max_steps,is_done",
+    [
+        (None, None, False),
+        (-1, None, False),
+        (None, -1, False),
+        (5, -1, False),
+        (-1, 10, False),
+        (None, 0, True),
+        (0, None, True),
+        (-1, 0, True),
+        (0, -1, True),
+    ],
+)
+def test_trainer_max_steps_and_epochs_fit_loop_done(max_epochs, max_steps, is_done):
+    trainer = Trainer(max_epochs=max_epochs, max_steps=max_steps)
+
+    assert trainer.max_epochs == max_epochs
+    assert trainer.max_steps == max_steps
+    assert trainer.max_time is None
+    assert trainer.fit_loop.done is is_done
+
 
 def test_trainer_min_steps_and_epochs(tmpdir):
     """Verify model trains according to specified min steps"""