How to release the CUDA Memory in torch hook function?

sev777: <p>I use a hook to store gradient information when I test, but this will cause CUDA out of memory, how do I release the information in the hook?</p>

ptrblck: <p>I assume storing the gradient information increases the memory usage and you are thus running OOM. You can delete tensors via <code>del tensor</code>.</p>

raining_day513: <p>Hi, ptrblck! Could you please help me with a little related question:</p>
<p>Is it possible to deallocate the CUDA memory earlier in the backward pass using hooks? It seems that PyTorch would do this at once for all gradients (I heard some people was saying this but I cannot confirm it myself)</p>
<p>Thanks in advance! Really thank you for consistently helping people on this forum!</p>

ptrblck: <p>I’m unsure what exactly you want to delete.</p>
<aside class="quote no-group" data-username="raining_day513" data-post="3" data-topic="163014">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/raining_day513/48/52553_2.png" class="avatar"> raining_day513:</div>
<blockquote>
<p>It seems that PyTorch would do this at once for all gradients</p>
</blockquote>
</aside>
<p><code>optimizer.zero_grad()</code> or <code>model.zero_grad()</code> will use <code>set_to_none=True</code> in recent PyTorch releases and will thus delete the <code>.grad</code> attributes of the corresponding parameters.<br>
However, this is done after calling <code>optimizer.step()</code> to update the parameters with the calculated gradients. Deleting gradients in a backward hook won’t make sense, since you won’t be able to update the parameters at all.</p>
<p>If you are thinking about the intermediate forward activations, I also don’t think you would be able to delete them in backward hooks, as I guess that references are still stored (outside of the backward hook). Even if you <code>del</code> the saved tensors or gradient tensors, I would guess Autograd would still keep the references and eventually free it when the last reference was deleted.<br>
However, I didn’t test this behavior so it’s a guess.</p>

raining_day513: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="163014">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Deleting gradients in a backward hook won’t make sense, since you won’t be able to update the parameters at all.</p>
</blockquote>
</aside>
<p>I would like to do it anyway with some personal experiments. So is there any API to remove the gradient of a given operator/backward node during the backward pass? (Or maybe you just answered me since the API is the <code>set_to_none</code> option?) (by remove I meant I want to release the GPU memory it has occupied during the backward pass. It’s OK if this requires writing some C++ or maybe I need to change some underlying codes)</p>
<p>Thanks for your reading.</p>

ptrblck: <p>I tried to delete the <code>.grad</code> attribute using <code>.register_hook</code>, but this didn’t work unfortunately.<br>
<a class="mention" href="/u/alband">@albanD</a> would you have a recommendation how to run such experiments?<br>
I know CPU-offloading uses context managers to at least move tensors to the CPU, but I assume these are forward activations (and other stored tensors needed for backward) and not the actual gradients?</p>

albanD: <p>Ho yes because such hook runs before we accumulate into the <code>.grad</code> field.<br>
I think you can either use autograd.grad to not have issues with .grad fields being populated.<br>
Another trick is to get a post-Node hook and clear the <code>.grad</code> field then.</p>
<p>An example I have where this is done is this one where I make optimizer a hook: <a href="https://gist.github.com/albanD/18c240bd2e09f9d93f5c4a0c9ccda39e" class="inline-onebox">PyTorch optimizer as hook · GitHub</a></p>

raining_day513: <p>Thanks for the example!</p>