Best way to proceed to find memory leaks during backprop?

Hamster: <p>Hello!</p>
<p>I’m currently having some CUDA <em>“out of memory”</em> errors during training after a certain number of iterations, and only when back-propagation is being used and precisely when trying to allocate new tensors in the <code>backward</code> methods of my <code>autograd::Function</code> subclasses. I’m using custom activation functions &amp; layers for which I wrote custom CUDA code which is then wrapped in <code>autograd::Function</code> subclasses. I am 99% sure it is due to memory leaks and not problem size as I did something pretty similar in Python/PyTorch and had no issues at all, and on my C++ code it happens only <strong>after a certain number of iterations</strong>.</p>
<p>So the thing is that I’d like to know how one should go in general about identifying memory leaks in a <code>libtorch</code> program because right now a direct use of say valgrind doesn’t help me much and debugging/guessing is currently a real pain. Say for example since, from my understanding, a caching CUDA allocator is used by libtorch, I’d like to find (for example during debugging) whether some tensors are indeed being freed (as in their space is marked as “available” by the caching allocator) after they’re being dereferenced, because if they aren’t then that can help me indeed to localize a bug in my code causing those tensors to not be properly dereferenced.</p>
<p>Thank you!</p>

ptrblck: <p>Since you are assuming your custom CUDA extension is causing the trouble, I would probably try to come up with a minimal standalone code snippet using your module only (forward and backward) and check the allocations via <code>print(torch.cuda.memory_summary())</code> in each iteration. If you are seeing a larger memory usage, then it would indeed point to your custom code. In this case, check first if you are storing any tensors, which would disallow PyTorch to free them. Also, I assume you are using <code>libtorch</code> or are you manually trying to manage the CUDA memory via raw malloc/free calls?</p>

Hamster: <p>Great idea, thanks a lot ptrblck! I’m working entirely in C++ though (both for training &amp; inference), so no <code>torch.cuda.memory_summary()</code> right now but by looking at the source code of the Python API I found that I can get what I need by playing around with <code>c10::cuda::CUDACachingAllocator::getDeviceStats</code> and the <code>c10::cuda::CUDACachingAllocator::StatType</code> enum.</p>

Hamster: <p>Found the culprit thanks to the memory stats. In a <code>torch::autograd::Function</code>, when I create a tensor during <code>forward</code> and store it in <code>ctx-&gt;saved_data</code> (for example <code>ctx-&gt;saved_data["some_tensor"] = some_tensor;</code> in <code>forward</code>, where <code>some_tensor</code> contains intermediary results I need to reuse later in the <code>backward</code> method), when we arrive at the <code>backward</code> method the tensor doesn’t get dereferenced even after <code>backward</code> finishes. However, if I add a <code>ctx-&gt;saved_data["some_tensor"] = torch::Tensor();</code> at the end of <code>backward</code> (explicitly dereferencing <code>some_tensor</code>), I don’t have any memory leaks anymore.</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> can you please confirm whether or not it is expected behavior that tensors inside <code>ctx</code> do <strong>not</strong> get dereferenced even after the end of <code>backward</code> ?</p>

ptrblck: <p>I think the issue is caused by the wrong usage of directly string the tensors in <code>ctx</code> and if I’m not mistaken, we’ve seen a few similar issues in the past posted here.<br>
Use <code>ctx-&gt;save_for_backward</code> and the memory leak should be gone.</p>

Hamster: <p>I confirm the problem is also gone when using <code>ctx-&gt;save_for_backward</code> instead of saving directly in the <code>saved_data</code> hash-map.</p>