Are the gradients of parameters in nn.Module all set to `None` and thus not allocated by default?

raining_day513: <p>I want to know when PyTorch will allocate GPU memory for all the gradients of a given <code>nn.Module</code>. It seems that all these gradients are set to <code>None</code> before training, and they will be set during the backward pass.</p>
<p>Setting to <code>None</code> or not, it seems that PyTorch already allocate GPU memory for all of them. Is this claim correct?</p>
<p>How can I change this behaviour so that I can decide when the GPU memory will be allocated for the gradients that are about to be computed.</p>

ptrblck: <aside class="quote no-group" data-username="raining_day513" data-post="1" data-topic="180081">
<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>Setting to <code>None</code> or not, it seems that PyTorch already allocate GPU memory for all of them. Is this claim correct?</p>
</blockquote>
</aside>
<p>No, that’s wrong as PyTorch will not allocate memory for non-existing objects.</p>

raining_day513: <p>I observed that there was not only one (weight) but three times (weight + gradient * 2) memory usage during forward, and my conclusion at that time was that it was caused by (1) the momentum mechanism of the optimizer that requires the gradients to be stored from the previous training loop (previous batch), and (2) the fact that PyTorch might have allocated all <code>.grad</code>s of a <code>nn.Module</code> after the instance is created. Based on your answer here, is my conclusion completely wrong? I also believe that this has something todo with the ordering of code I mentioned in <a href="https://discuss.pytorch.org/t/does-an-optimizer-deallocate-the-gradients-after-it-updates-the-weights-of-model/180084/3">another thread yesterday</a>.</p>
<p>Could you help me confirm this, as I might need to report this back to my teammate?</p>
<p>Thanks for your reading.</p>

ptrblck: <aside class="quote no-group" data-username="raining_day513" data-post="3" data-topic="180081">
<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>Based on your answer here, is my conclusion completely wrong?</p>
</blockquote>
</aside>
<p>Yes, since <code>None</code> objects are not taking up GPU memory.<br>
Of course if you <em>do not delete</em> the <code>.grad</code>attributes they will still be in memory and the peak memory usage will increase as described in the other thread.<br>
However, here you claim that PyTorch allocates memory for gradients even if these are non-existent, which is wrong.</p>
<p>You would also have to check the size of the intermediate tensors created during the forward pass, which would also use device memory until they are freed during the backward pass.</p>

raining_day513: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="180081">
<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>However, here you claim that PyTorch allocates memory for gradients even if these are non-existent, which is wrong.</p>
</blockquote>
</aside>
<p>Firstly, apologize for the wrong claim. I’m learning this. I was not sure whether these gradients exist or not after the model instance is created.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="180081">
<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>You would also have to check the size of the intermediate tensors created during the forward pass, which would also use device memory until they are freed during the backward pass.</p>
</blockquote>
</aside>
<p>I didn’t mention that I implemented the gradient-checkpointing technique(or activation checkpointing since the word checkpoint has other meanings) using <code>torch.autograd.Function</code>. By this, I think the intermediate tensors are not involved in the problem.</p>

ptrblck: <aside class="quote no-group" data-username="raining_day513" data-post="5" data-topic="180081">
<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>I was not sure whether these gradients exist or not after the model instance is created.</p>
</blockquote>
</aside>
<p>They don’t and you can check if my accessing the <code>.grad</code> attribute of any parameter which will return <code>None</code>.</p>
<aside class="quote no-group" data-username="raining_day513" data-post="5" data-topic="180081">
<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>I didn’t mention that I implemented the gradient-checkpointing technique(or activation checkpointing since the word checkpoint has other meanings) using <code>torch.autograd.Function</code>. By this, I think the intermediate tensors are not involved in the problem.</p>
</blockquote>
</aside>
<p>Intermediate activation tensors are stored by default during the forward pass as they are needed to calculate the gradients in the backward pass. Activation checkpointing would recompute these activations to save memory and since you aren’t using it Autograd will still store the intermediates thus increasing the memory usage.</p>

raining_day513: <aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="180081">
<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>Activation checkpointing would recompute these activations to save memory and since you aren’t using it Autograd will still store the intermediates thus increasing the memory usage.</p>
</blockquote>
</aside>
<p>Could you elaborate more on this part? Did you mean that the intermediate data will still be stored even if I put the computation inside the block of <code>torch.no_grad()</code> when defining the static method <code>backward</code> of <code>torch.autograd.Function</code>? I did implement activation checkpointing in the project.</p>

ptrblck: <p>No, <code>no_grad</code> will prevent storing intermediates, but in case you are using custom <code>autograd.Function</code>s you would need to explicitly store the activations via <code>ctx.save_for_backward</code>. Otherwise you won’t be able to compute gradients.</p>

raining_day513: <aside class="quote no-group" data-username="raining_day513" data-post="5" data-topic="180081">
<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>I didn’t mention that I implemented the gradient-checkpointing technique</p>
</blockquote>
</aside>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> Lol, I think my bad grammar caused the confusion. I hope you haven’t read my reply above(already deleted). I was trying to say:</p>
<ol>
<li>I did implement gradient-checkpointing so the issue should not be related to intermediate data in my case.</li>
<li>I forgot to mention point 1.</li>
</ol>
<p>I just realized that my original sentence sounds like:</p>
<ol>
<li>I didn’t implement something.</li>
<li>The intermediate data was not involved because of 1.</li>
</ol>
<p>Really sorry for this if it causes confusion!</p>
<p>Back to the title, if I’m very sure that I have implemented gradient-checkpointing but I still saw some abnormal GPU memory caching(from <code>nvidia-smi</code>) that is much larger than what I calculate for my model, what could be the cause? This is why I made some (probably wrong) assumptions that PyTorch didn’t allocate the GPU memory efficiently.</p>