Lightning-AI / Lightning-AI/pytorch-lightning
Enable batch size finder for distributed strategies
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
It's not clear why it's currently disabled [here](https://github.com/Lightning-AI/pytorch-lightning/blob/c235f20e7131af2c7be4cc9080d3c946d93d58ea/src/lightning/pytorch/callbacks/batch_size_finder.py#L137).
### Pitch
There should not be a big difference in how it works vs. LR finder. E.g. all ranks try the same size under try/catch and reduce 1 or 0 based on whether it was successful. This operation is repeated with a given strategy until all ranks are successful.
### Alternatives
Manual HPO
### Additional context
_No response_
cc @borda
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/lightning/pytorch/callbacks/batch_size_finder.py at the disabled logic around line 137, then compare the existing LR finder behavior. Trace how distributed strategies coordinate ranks and handle failures. Done means batch size finding works across distributed strategies, with all ranks testing the same size and agreeing on success before continuing.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100