scale_batch_size runs zero batches per trial when started after training has begun #21903

Bug description

scale_batch_size runs zero training batches per trial when the search starts after training has already taken some optimizer steps. Nothing ever raises an OOM, so every trial "succeeds", the search keeps doubling, and it returns a batch size that the model cannot actually run. Training then dies with the OOM the finder exists to prevent.

This is the pattern the BatchSizeFinder docstring itself documents, calling scale_batch_size from on_train_epoch_start at a set of milestones:

class FineTuneBatchSizeFinder(BatchSizeFinder):
    def on_fit_start(self, *args, **kwargs):
        return

    def on_train_epoch_start(self, trainer, pl_module):
        if trainer.current_epoch in self.milestones or trainer.current_epoch == 0:
            self.scale_batch_size(trainer, pl_module)

There are two independent causes, which is why I am filing this rather than sending a one-line patch.

1. max_steps is absolute, but is set as if it were relative.

__scale_batch_reset_params (src/lightning/pytorch/tuner/batch_size_scaling.py:146) does:

trainer.fit_loop.epoch_loop.max_steps = steps_per_trial

_FitLoop.done compares that against the absolute global_step (fit_loop.py:179). _try_loop_run restores the dumped loop state before every trial, so global_step is whatever it was when the search started. Once global_step >= steps_per_trial, done is True immediately and the trial is a no-op. The sibling tuner gets this right: lr_finder.py:338 uses num_training + trainer.global_step.

2. The restored batch_progress puts each trial partway through an epoch.

The loop state dumped at search start carries batch_progress.current.ready from the epoch already in flight. Every trial restores it, so a trial whose dataloader has fewer batches than that value finishes instantly. Adding the global_step offset from cause 1 is not enough on its own.

Instrumented trials, steps_per_trial=3, max_trials=4, model OOMs above batch_size=8. Search at epoch 0 works; search at epoch 1 (global_step=5) runs nothing:

>>> search at epoch=0 global_step=0 batch_progress.current.ready=0
    bs=  2 num_train_batches=   5 batch_ready_after_restore=  3 steps_run= 3
    bs=  4 num_train_batches=  16 batch_ready_after_restore=  3 steps_run= 3
    bs=  8 num_train_batches=   8 batch_ready_after_restore=  3 steps_run= 3

>>> search at epoch=1 global_step=5 batch_progress.current.ready=5
    bs=  2 num_train_batches=   5 batch_ready_after_restore=  5 steps_run= 0
    bs=  4 num_train_batches=  16 batch_ready_after_restore=  5 steps_run= 0
    bs=  8 num_train_batches=   8 batch_ready_after_restore=  5 steps_run= 0
    bs= 16 num_train_batches=   4 batch_ready_after_restore=  5 steps_run= 0

searches (global_step, found_batch_size): [(0, 8), (5, 16)]

Applying only the + trainer.global_step offset fixes cause 1 but not cause 2, so the search still returns 16:

>>> search at epoch=1 global_step=5 batch_progress.current.ready=5
    bs=  2 num_train_batches=   5 batch_ready_after_restore=  5 steps_run= 0
    bs=  4 num_train_batches=  16 batch_ready_after_restore=  8 steps_run= 3
    bs=  8 num_train_batches=   8 batch_ready_after_restore=  8 steps_run= 3
    bs= 16 num_train_batches=   4 batch_ready_after_restore=  5 steps_run= 0

searches (global_step, found_batch_size): [(0, 8), (5, 16)]

The bs=16 trial has 4 batches but starts at batch 5, so it still never runs and never OOMs.

A correct fix needs the trials to start from a clean epoch position while still returning to the real position afterwards. _try_loop_run currently uses one dumped state for both purposes. That is a design decision I did not want to make unilaterally, so I am reporting it instead of guessing.

What version are you seeing the problem on?

master

How to reproduce the bug

from unittest.mock import patch

from torch.utils.data import DataLoader

from lightning.pytorch import Trainer
from lightning.pytorch.callbacks.batch_size_finder import BatchSizeFinder
from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset

OOM_ABOVE = 8
found = []


class Model(BoringModel):
    searching = False

    def __init__(self):
        super().__init__()
        self.batch_size = 2

    def training_step(self, *a, **k):
        if self.searching and self.batch_size > OOM_ABOVE:
            raise RuntimeError("CUDA error: out of memory")
        return super().training_step(*a, **k)

    def train_dataloader(self):
        return DataLoader(RandomDataset(32, 64), batch_size=self.batch_size)


class FineTuneBatchSizeFinder(BatchSizeFinder):
    def on_fit_start(self, *args, **kwargs):
        return

    def on_train_epoch_start(self, trainer, pl_module):
        pl_module.searching = True
        self.scale_batch_size(trainer, pl_module)
        pl_module.searching = False
        found.append((trainer.global_step, pl_module.batch_size))


@patch("lightning.pytorch.tuner.batch_size_scaling.is_oom_error", return_value=True)
def main(_):
    Trainer(
        max_epochs=2,
        limit_train_batches=5,
        limit_val_batches=0,
        num_sanity_val_steps=0,
        enable_progress_bar=False,
        enable_model_summary=False,
        logger=False,
        enable_checkpointing=False,
        callbacks=[FineTuneBatchSizeFinder(steps_per_trial=3, max_trials=4)],
    ).fit(Model())
    print("(global_step at search, batch size found):", found)


main()

Error messages and logs

(global_step at search, batch size found): [(0, 8), (5, 16)]

Both searches see the same model, so they should agree. The second returns 16, which does not fit. Without the searching guard in the snippet above, the run instead ends in RuntimeError: CUDA error: out of memory right after the second search.

Environment

#- PyTorch Lightning Version: 2.6.2 (master, fcef40451)
#- PyTorch Version: 2.13.0
#- Python version: 3.11
#- OS: macOS 15 (arm64)
#- CUDA/cuDNN version: none, CPU
#- How you installed Lightning: source, editable

More info

All 54 tests in tests/tests_pytorch/tuner/test_scale_batch_size.py start the search at global_step == 0, either through on_fit_start or through tuner.scale_batch_size before fit, so none of them reach this path.