Member-only story
Mastering GPU Memory Management With PyTorch and CUDA
A gentle introduction to memory management using PyTorch’s CUDA Caching Allocator
My articles are free to read for everyone. If you don’t have a Medium subscription, read this article by following this link.
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00
MiB. GPU 0 has a total capacity of 79.32 GiB of which 401.56 MiB is free.If you’ve ever tried training large deep learning models in PyTorch, or have worked with larger datasets, the above would seem very familiar to you. The dreaded CUDA out of memory error! This happens whenever your GPU runs out of memory while trying to allocate space for tensors and can be very frustrating to deal with, especially when you’ve spent a lot of time fine-tuning your model and optimizing your code.
I’ve previously shared some insights on how you can train larger deep learning models in PyTorch — much faster and with lower memory usage. If you haven’t checked them out already, I highly recommend giving them a read as this article will build upon that.
The strategies described in these articles do help alleviate this issue to some extent, however, they only cover techniques for training speed and efficiency. In this article, we’ll take a deep dive into how PyTorch optimizes GPU memory usage, and how you can tailor some of it’s internal systems to squeeze that extra oomph out of your GPU cluster.
The Importance of GPU Memory Management
In today’s world of exponentially growing datasets and increasingly sophisticated models, the efficient use of GPU memory has become a top priority. No matter how powerful your GPU, the…
