Add BatchSizeFinder callback

[rohitgr7]

What does this PR do?

Part of #11012

TODO:

• Add docs in followups
• Address added TODOs here in the followup

Does your PR introduce any breaking changes? If yes, please list them.

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

cc @Borda @awaelchli @ananthsub @daniellepintz @rohitgr7 @akihironitta

[tchaton]

After chatting with @rohitgr7,

Here is the API design specification for the BatchSizeFinder.

With the tune method.

User story: As a user, I want to fit the best batch size for each stage and get the associated maximal batch size.

trainer = Trainer()

trainer.tune(auto_batch_size="fit | True (BC)", ...) # Compute the train_batch_size and val_batch_size
   # In the LightningModule, DataModule
    self.batch_size = train_batch_size # BC -> depreceated in favor of `{stage}_batch_size`.
    self.train_batch_size = train_max_batch_size
    self.val_batch_size = val_max_batch_size

trainer.tune(auto_batch_size='validate', ...) -> val dataloader batch_size
    # In the LightningModule, DataModule
    self.val_batch_size = max_batch_size

trainer.tune(auto_batch_size='test', ...) -> test dataloader batch_size
    # In the LightningModule, DataModule
    self.test_batch_size = max_batch_size

trainer.tune(auto_batch_size='predict', ...) -> predict dataloader batch_size
    self.predict_batch_size = max_batch_size

With the Trainer entry point method.

User story: As a user, I want to fit the best batch size and run the trainer entry point into a single call.

trainer = Trainer(callbacks=BatchSizeFinder())
trainer.fit(...) -> `on_fit_start` perform tune
trainer.validate(...) -> `on_validation_start` perform tune
trainer.test(...) -> `on_test_start` perform tune
trainer.predict(...) -> `on_predict_start` perform tune

Open questions:

1. Some use case requires dynamic maximum batch size such as fine-tuning where the memory footprint of the model increases during trainer.

Solution:

1. Enable a scheduling strategy. Would this be enough to cover use cases?
2. Re-compute the max_batch_size if the params group of the optimizer changed?

...

@carmocca @ananthsub @awaelchli @justusschock thoughts ?

[rohitgr7]

update on PR:

all that's been discussed here is now added.

Only the following proposal is missing and I need more time to think about how to go about this. Maybe in a sep PR to avoid a lot of changes in a single one.

trainer.tune(auto_batch_size="fit | True (BC)", ...) # Compute the train_batch_size and val_batch_size
   # In the LightningModule, DataModule
    self.batch_size = train_batch_size # BC -> depreceated in favor of `{stage}_batch_size`.
    self.train_batch_size = train_max_batch_size
    self.val_batch_size = val_max_batch_size

Also the API is:

trainer.tune(method="fit|validate|test|predict") (default fit)

instead of

trainer.tune(auto_batch_size="fit|validate|test|predict")

because, there is another option to run batch_size_finder using trainer.tuner.scale_batch_size() and auto_batch_size argument won't be a good one here. Also auto_batch_size is a Trainer argument already.

[ananthsub]

IMO a callback is not the right fit here. this will override the behavior of fit/validate/test if used this way. why doesn't this use the loops API instead?

[rohitgr7]

IMO a callback is not the right fit here. this will override the behavior of fit/validate/test if used this way. why doesn't this use the loops API instead?

I tried writing pseudocode to cover up all the possible use-cases we have and convert it into a component that users can easily integrate and plugin, but loops couldn't fit all possible cases. For eg, we have flash finetuning and let's say the user might want to compute new LR or new batch_size after certain epochs of pre-training, then it's not easily configurable within a single call. One can achieve it with multiple calls but since we support strategies within Flash. With callbacks this is now pretty simple since we make sure after each tuning call, the respective loop state_dict is restored completely, and we still get an updated hparam using the tuner. Plus it now runs for val/test/predict.

If you have any ideas to configure it using loops, would love to discuss them.

Also, we discussed the possible solutions for this internally and everyone I talked to recommended callbacks.
#9103
#11012 (comment)

[carmocca]

Synced outside of GitHub! LGTM! Great and challenging work

We can try to remove the exception in the future if we finish #8604