Why CUDA extensions use more gpu memory?

nobody: <p>I used custom cuda extension to replace some parts of the model, and it worked correctly as expected. But training with cuda extension uses more GPU memory.<br>
I used <code>torch.cuda</code> to analyze the pre-expansion and post-expansion pytorch GPU memory allocation. I found that in addition to the GPU memory I pre-allocated for cuda extensions using <code>cuda_malloc</code>, pytorch itself allocates more GPU memory.<br>
Half of the over-allocated GPU memory is due to the explicit call <code>torch.empty_like()</code> in the <code>forward</code> and <code>backward</code> of the cuda extension to pre-allocate the space for storing the calculation results, and the other half is added when the <code>optimizer.step()</code> is called.<br>
I need some help information to locate the problem and solve the problem of using the video memory.</p>
<ol>
<li>Does the GPU memory statistics method (<code>memory_allocated</code> or <code>memory_summary</code>) in <code>torch.cuda</code> count the GPU memory allocated by me manually calling the <code>cuda_malloc</code> function? Do it count the GPU memory allocated by similar methods called <code>torch::empty_like()</code> in C++ code? My guess is that the former is not counted, but the latter is counted.</li>
<li>Will the GPU memory allocated by <code>torch::empty_like()</code> called in the custom forward and backward methods in the c++ code be automatically recycled when it is no longer used? If not, then how should I destory this part of the GPU memory by myself?</li>
<li>Why does the call of <code>optimizer.step()</code> lead to the increase of the GPU memory usage? I think this method is only to update the weights and should not allocate new space, but it does increase the GPU memory usage.</li>
</ol>

tom: <ol>
<li>This seems as you describe.</li>
<li>Memory is returned to the torch cache once the tensors using it go out of scope.</li>
<li>It depends on your optimizer, but things like tracking statistics for adaptivity (eg in Adam) or momentum use additional memory. These are allocated once, so it should not increase over time.</li>
</ol>
<p>In general, I would advise to no use <code>cuda_malloc</code> directly but instead create a new torch tensor. As long as you keep a reference to it, the memory will be reserved for you and can be accessed via <code>data_ptr()</code> (for a void*) or <code>data_ptr&lt;T&gt;()</code> for a <code>T*</code>.</p>
<p>Best regards</p>
<p>Thomas</p>

nobody: <p>Regarding the third point, maybe my description above is not clear enough.</p>
<p>I compared the GPU memory usage at each step of the code that does not use the CUDA extension and the code that uses the CUDA extension. Both use the exact same optimizer. My cuda extension only replaces the layer modules in the model.</p>
<p>The former only adds about 5M of GPU memory usage during the first execution of <code>optimizer.step()</code>, while the latter adds about 1G of GPU memory usage. There is a huge difference between the two.</p>
<p>I don’t know why the CUDA extension causes such a difference? Can you provide some possible clues?</p>

tom: <p>I would not know why there would be such a difference. Might be something your extension does?</p>

nobody: <p>Thank you for your answer above. I found the reason. The optimizer uses different gradient calculation strategies for one-dimensional parameters, which requires more GPU memory space.</p>