yoyodyne icon indicating copy to clipboard operation
yoyodyne copied to clipboard

Automatic batch sizing

Open kylebgorman opened this issue 1 year ago • 9 comments

PyTorch Lightning supports a mode where it automatically computes the maximum batch size on your accelerator (by collecting gradients over a few batches and then using binary search to find a max that doesn't give OOM errors). Documentation is here. It appears straightforward to implement (see "Note" on docs.)

Having this enabled would be very useful for computing the maximum batch size possible. From this one could factor the desired batch size and accumulate gradients across multiple batches per optimizer step, so this is related to #132 in the obvious way.

This will also combine nicely with a move to LightningCLI, IMO, since that uses a subcommand interface and we could do something like yoyodyne scale_batch_size --arg1 --arg2 to enable this.

kylebgorman avatar Nov 01 '23 21:11 kylebgorman