Add support for handling training/validation OOMs gracefully

a-r-j / ProteinWorkshop

Benchmarking framework for protein representation learning. Includes a large number of pre-training and downstream task datasets, models and training/task utilities. (ICLR 2024)

https://proteins.sh/

MIT License

194 stars 16 forks source link

Add support for handling training/validation OOMs gracefully #81

Closed amorehead closed 7 months ago

amorehead commented 7 months ago

Adds single-GPU support for handling training/validation OOMs gracefully.
Note that multi-GPU support for OOM-catching (although in principle handled by this code) does not currently work and may lead to unpredictable stalls during training/validation forward (or backward) passes (with the latest release of Lightning - note that this behavior may change with future updates to Lightning).

amorehead commented 7 months ago

Sure thing. This issue for Lightning seems most relevant and is actually where I got this initial implementation: https://github.com/Lightning-AI/pytorch-lightning/issues/5243#issuecomment-1553404417