Lightning-AI / Lightning-AI/pytorch-lightning

BatchSizeFinder throws KeyError: 'limit_eval_batches'

Open
#18,985 4 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

bug duplicate help wanted repro needed tuner ver: 2.1.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

Using the latest Lightning-AI v 2.1

I have added vanilla BatchSizeFinder as a callback to my lightningmodule. After it finishes finding the correct batchsize, it breaks while calling function batch_size_scaling. The error message is:
```
File ~/miniconda3/envs/dl/lib/python3.10/site-packages/lightning/pytorch/tuner/batch_size_scaling.py:155, in __scale_batch_restore_params(trainer, params)
153 stage = trainer.state.stage
154 assert stage is not None
--> 155 setattr(trainer, f"limit_{stage.dataloader_prefix}_batches", params["limit_eval_batches"])
157 loop.load_state_dict(deepcopy(params["loop_state_dict"]))
158 loop.restarting = False
```
On debugging the params variable only has the following keys:
```
dict_keys(['loggers', 'callbacks', 'max_steps', 'limit_train_batches', 'limit_val_batches', 'loop_state_dict'])
```
So indeed the limit_eval_batches key does not exist in my lightningmodule. Happy to provide more info if needed.

Thanks

### What version are you seeing the problem on?

v2.1

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

_No response_

### More info

_No response_

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

Inspect lightning/pytorch/tuner/batch_size_scaling.py, starting at __scale_batch_restore_params where the traceback accesses limit_eval_batches. Reproduce the BatchSizeFinder callback flow and verify that scaling completes when the parameters contain only the listed train and validation limits, without raising KeyError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.