Automatic gradient accumulation for batch size tuning

Description & Motivation

When training different model sizes on a different number of devices or different hardware, the batch size needs to be carefully tuned in order to achieve maximum GPU utilization without incurring Out Of Memory Error. Lightning currently has a BatchSizeFinder tuner, but it does not work with distributed training, which is precisely the setup where tuning the batch size is a tricky and expensive process.

Pitch

One appealing solution could be to use an automatic microbatching strategy with gradient accumulation. In terms of user API it could be used with

Trainer(
    ...
    accumulate_grad_batches="auto",
    ...
)

The strategy is the following: we attempt to train on the full batch size provided by the user, catch any Cuda OOM exception, half the batch size if needed and retry training until the forward step succeeds. These micro batches are then accumulated to simulate the effective batch size.

(Note that this does mean that we should be able to slice through the batches provided by the user which can be challenging if the user uses custom data structures or layout, but should be doable in the most standard practice)

This feature maximizes efficiency while making it easier for the user to scale in and out without having to constantly tune the batch size.

Alternatives

If you have any recommendation or alternatives for automating the batch size in distributed training, I would be very thankful for that!

Additional context

I came across this feature when reading through the documentation of the fantastic Mosaic ML library Automatic Microbatching and found the idea very cool and useful. It would be a very nice addition to Lightning as there are no way of automatically tuning the batch size in distributed settings afaik.
Their implementation is open source, under Apache 2.0.

cc @borda

Start by reviewing the existing BatchSizeFinder and Trainer handling for accumulate_grad_batches, then compare the linked Mosaic ML automatic microbatching implementation. Define how distributed retries, batch slicing, and custom data structures are handled; done means automatic accumulation works for the described distributed training cases without out-of-memory failures.