Lightning-AI / Lightning-AI/pytorch-lightning

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

Open
#21,903 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

ver: 2.7.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### 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:

```python
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:

```python
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

```python
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

Current 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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in src/lightning/pytorch/tuner/batch_size_scaling.py, especially __scale_batch_reset_params and _try_loop_run, then compare the offset handling in lr_finder.py:338 and completion logic in fit_loop.py:179. Run tests/tests_pytorch/tuner/test_scale_batch_size.py and extend coverage for searches begun after training has advanced. Done means trials execute the requested steps from a clean epoch position, restore the real loop state afterward, and detect the OOM boundary consistently.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.