Support BatchSizeFinder in DDP

[AlJohri]

Description & Motivation

Based on the docs, BatchSizeFinder does not currently support DDP:

Batch size finder is not yet supported for DDP or any of its variations, it is coming soon.
https://lightning.ai/docs/pytorch/stable/advanced/training_tricks.html

It would be great to have a way to automatically find the correct batch size when using submitting DDP jobs to a queue that has multiple instance types. We would like to automatically detect the correct batch size to use when doing single instance, multi-gpu training.

Pitch

Support BatchSizeFinder when using DDP strategy.

Alternatives

Alternatively, provide some way to run BatchSizeFinder on a single GPU prior to DDP being used.

Additional context

Related Issues and PRs:

• Tuner fails with DDP in 1.5 #10560
• Update docs regarding support of tuner with DDP #10754
• Add BatchSizeFinder callback #11089
• Disable tuner with distributed strategies #12179
• Support DDP for BatchSizeFinder #15305

cc @lantiga @Borda @justusschock

[davidgill97]

I dont think there is an easy way around this: DDP creates multiple processes, and the current impl of batchsizefinder increases batch size until cuda OOM occurs, and backtracks (if i understood it correctly), which requires creating multiple processes every trial.
The current method i use for single node multi gpu setup is: 1. find gpu with smallest memory (can be skipped if your gpus are same), 2. run bs find on one gpu, 3. use the bs globally.

TBH, current batchsizefinder is very inefficient.. Better heuristic i adopted was to 1. set bs=1, 2. run for single epoch and measure peak memory, 3. deduct model size, expected gradient size, etc that do not scale with bs, 4. then just divide gpu memory by the memory usage.