Batch size should be automatically computed

[ohinds]

Batch size should be automatically computed

[hvgazula]

I recommend using accelerate. For more information, please refer to https://huggingface.co/docs/accelerate/usage_guides/memory or https://towardsdatascience.com/a-batch-too-large-finding-the-batch-size-that-fits-on-gpus-aef70902a9f1. I personally prefer the former as I find it to be much cleaner.

[hvgazula]

Unfortunately, accelerate only supports pytorch. Probably, will have to wait until tensorflow is supported.

[hvgazula]

Unfortunately, tensorflow doesn't have decorators/functions to auto-scale batch size like how lightning/accelerate for pytorch does.

However, here's a naive example of accomplishing this.

@satra let me know your thoughts about where in the codebase should this go, API or docs.

[satra]

perhaps add that as a more generic utility function, and the api can use an enum like AUTO_TUNE_BATCH to indicate whether that function should be called or the provided batch size used. however, as we know this is not always a function of largest batch size possible given memory requirements. hence we should consider situations where we don't need all the memory and could even request shards corresponding to the more optimal batch size.