Help with CUDA memory allocation during forward Linear

Monkey.Py: <p>Hi I am currently working on some profiling of my model and I have some question about memory allocation with torch.<br>
The code I am using is :</p>
<pre><code class="lang-auto">device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
torch.cuda.memory._record_memory_history(True,True,1000,True,device,True)
xQuery = torch.randn(1,259, 259,128).to(device)
xFocalMaps=[torch.randn(1, 259, 259, 128).to(device), torch.randn(1, 130, 130, 128).to(device), torch.randn(1,65, 65, 128).to(device)]
qkv=self.qkv(xQuery) # LINEAR(123,3*128)
s=torch.cuda.memory._snapshot()
with open("MySnapshot.pickle", "wb") as f:
     pickle.dump(s, f)
</code></pre>
<p>I will describe it quickly<br>
 → XQuery a float32 tensor (1,259,259,128) so allocated memory is 4*128*259*259 / 1024**2 = 32,8Mb.<br>
 → XFocalMap 3 float 32 tensors of different shapes (I will give direct size result) of size [ 32,8Mb ; 8,3Mb ; 2,1Mb ].<br>
 → self.qkv is Linear layer (as it is already initialized it is already allocated in memory)<br>
 → thus self.qkv(XQuery) will give a (1,259,259,3*128) float32 tensor = 98.3Mb.<br>
I use preforward and postforward hooks that print the cuda allocated memory in console, and I use memory snapshot as you can see in the code.</p>
<pre><code class="lang-auto">   layer_idx  call_idx layer_type       exp hook_type     mem_all  mem_cached  mem_all_diff  mem_cached_diff
0          1         0     Linear  baseline       pre   81.884766        90.0      0.000000              0.0
1          1         1     Linear  baseline       fwd  188.272949       210.0    106.388184            120.0
</code></pre>
<p>This is the result given by the hooks preforward and post forward of self.qkv(XQuery) : the difference between Pre and post allocated cuda memory is 106,4Mb &gt; 98.3Mb.<br>
This means that the result of self.qkv(XQuery) produces a 98.3Mb tensor and… something else which is 8.1Mb.<br>
Using the snapshot of the memory I get the following allocation scheme :</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/1/115eb1d52f16345fccd494d3bbfc0d68631269cf.png" data-download-href="https://discuss.pytorch.org/uploads/default/115eb1d52f16345fccd494d3bbfc0d68631269cf" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/115eb1d52f16345fccd494d3bbfc0d68631269cf_2_690x271.png" alt="image" data-base62-sha1="2tEZXzFGIcecvWLIInEX8pK0kKr" width="690" height="271" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/115eb1d52f16345fccd494d3bbfc0d68631269cf_2_690x271.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/115eb1d52f16345fccd494d3bbfc0d68631269cf_2_1035x406.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/1/115eb1d52f16345fccd494d3bbfc0d68631269cf.png 2x" data-dominant-color="E4E0DD"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1295×509 73.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>This picture shows the allocation of each XQuery, self.qkv(XQuery), XFocalMaps and my “somethingElseObject” of 8.1Mb in orange.<br>
Do you have an idea of what it could be ? What else other than the output is saved in the memory during the forward operation ?<br>
Thank you for help</p>

ptrblck: <p>The default cuBLAS workspace size for sm&lt;90 uses 8.125MB and is initialized <a href="https://github.com/pytorch/pytorch/blob/1460e5b7f5e85de532d18ad0c2ac7d48211100e8/aten/src/ATen/cuda/CublasHandlePool.cpp#L51C56-L51C87">here</a>:</p>
<pre><code class="lang-python">(4096 * 1024 * 2 + 16 * 1024 * 8) / 1024**2
8.125
</code></pre>

Monkey.Py: <p>Thank you for your help !!</p>

MoJo494: <p>Hey,</p>
<p>I have a follow-up question regarding the cuBLAS workspace size. When I observe GPU memory usage during fine-tuning a BERT-large model, I Initially observe the expected 8.125 MiB allocation for the cuBLAS workspace on my GPU. However, after completing the first forward pass in the fine-tuning process, I noticed an additional 8.125 MiB being allocated, resulting in a total of 16.25 MiB for cuBLAS workspace memory usage.</p>
<p>Could this indicate that the cuBLAS workspace size is dynamically adjusted during the fine-tuning process? If so, what factors might contribute to this increase in the cuBLAS workspace size? Your insights on this would be greatly appreciated. Thank you for your assistance!</p>

ptrblck: <p>PyTorch uses thread-local cuBLAS handles and creates a separate thread for the backward pass, which will thus allocate the workspace again for the backward cuBLAS handle.</p>

MoJo494: <p>Thank you for your swift and informative reply!</p>