feat: warn when training batches are fewer than accumulate_grad_batches

On-behalf-of: @Bocconi-Open-Source-Society

Author: giacomo-ciro

src/lightning/pytorch/CHANGELOG.md

@@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `log_key_prefix` parameter to `LearningRateMonitor` callback for prefixing logged metric names ([#21612](https://github.com/Lightning-AI/pytorch-lightning/issues/21612))
 
+- Added a `PossibleUserWarning` when the number of training batches is smaller than `accumulate_grad_batches`, since Lightning always steps on the last batch of the epoch and gradients end up accumulated over fewer batches than configured ([#21808](https://github.com/Lightning-AI/pytorch-lightning/issues/21808))
+
 ### Changed
 
 -

src/lightning/pytorch/loops/fit_loop.py

@@ -325,6 +325,15 @@ def setup_data(self) -> None:
                 category=PossibleUserWarning,
             )
 
+        if self.max_batches < trainer.accumulate_grad_batches:
+            rank_zero_warn(
+                f"The number of training batches ({self.max_batches}) is smaller than "
+                f"`accumulate_grad_batches={trainer.accumulate_grad_batches}`. Since Lightning always performs an"
+                " optimizer step on the last batch of the epoch, gradients will effectively be accumulated over only"
+                f" {self.max_batches} batch(es) instead of {trainer.accumulate_grad_batches} for this epoch.",
+                category=PossibleUserWarning,
+            )
+
     @property
     def restarted_on_epoch_start(self) -> bool:
         return self._restart_stage == RestartStage.RESTARTED_ON_EPOCH_START

tests/tests_pytorch/trainer/test_dataloaders.py

@@ -690,6 +690,22 @@ def test_warning_with_small_dataloader_and_logging_interval(log_interval, tmp_pa
         trainer.fit(model)
 
 
+def test_warning_with_dataloader_shorter_than_accumulate_grad_batches(tmp_path):
+    """Test that a warning message is shown if the dataloader has fewer batches than `accumulate_grad_batches`, since
+    Lightning always steps on the last batch of the epoch and gradients end up accumulated over fewer batches than
+    configured."""
+    model = BoringModel()
+    dataloader = DataLoader(RandomDataset(32, length=2))
+    model.train_dataloader = lambda: dataloader
+
+    trainer = Trainer(default_root_dir=tmp_path, max_epochs=1, accumulate_grad_batches=4, logger=False)
+    with pytest.warns(
+        UserWarning,
+        match=r"The number of training batches \(2\) is smaller than `accumulate_grad_batches=4`",
+    ):
+        trainer.fit(model)
+
+
 def test_warning_with_small_dataloader_and_fast_dev_run(tmp_path):
     """Test that a warning message is shown if the dataloader length is too short for the chosen logging interval."""
     model = BoringModel()