Gpu memory leak on torch.cat

Qmin-Lee: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/f/9f1cd22123faa855bc2034e70d3fe504e2aff0af.png" data-download-href="https://discuss.pytorch.org/uploads/default/9f1cd22123faa855bc2034e70d3fe504e2aff0af" title="SmartSelectImage_2022-06-07-14-34-25"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/f/9f1cd22123faa855bc2034e70d3fe504e2aff0af_2_690x488.png" alt="SmartSelectImage_2022-06-07-14-34-25" data-base62-sha1="mHzFiyWIWNSXAsbYSrxQse1xaGj" width="690" height="488" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/f/9f1cd22123faa855bc2034e70d3fe504e2aff0af_2_690x488.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/f/9f1cd22123faa855bc2034e70d3fe504e2aff0af_2_1035x732.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/f/9f1cd22123faa855bc2034e70d3fe504e2aff0af.png 2x" data-dominant-color="232423"><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">SmartSelectImage_2022-06-07-14-34-25</span><span class="informations">1105×783 69.1 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>
<hr>
<pre><code>                                               Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  
</code></pre>
<hr>
<pre><code>                                      aten::resize_         0.10%      69.000us         0.10%      69.000us       8.625us           0 b           0 b       4.74 Gb       4.74 Gb             8  
                                          aten::cat         0.04%      27.000us         0.22%     159.000us      39.750us           0 b           0 b       4.73 Gb           0 b             4  
                                         aten::_cat         0.08%      60.000us         0.19%     132.000us      33.000us           0 b           0 b       4.73 Gb           0 b             4  
                               attention_embedding2         0.13%      90.000us         0.74%     523.000us     523.000us          -4 b          -4 b       4.71 Gb           0 b             1  
                                          aten::mul         0.18%     129.000us         0.23%     162.000us      23.143us           0 b           0 b       1.57 Gb       1.57 Gb             7  
                               attention_embedding1         0.08%      54.000us         0.16%     112.000us     112.000us          -4 b          -4 b       1.56 Gb           0 b             1  
                                        aten::empty         0.70%     498.000us         0.70%     498.000us       4.835us       2.19 Kb       2.19 Kb     235.99 Mb     235.99 Mb           103  
                                         aten::lstm         0.15%     109.000us        92.66%      65.808ms      16.452ms           0 b           0 b      94.76 Mb           0 b             4  
                                   aten::_cudnn_rnn        35.10%      24.929ms        92.45%      65.656ms      16.414ms           0 b           0 b      94.76 Mb    -135.34 Mb             4  
                                   output_embedding         0.54%     380.000us        14.40%      10.229ms      10.229ms          -4 b        -372 b      58.78 Mb     -31.27 Mb             1  
</code></pre>
<hr>
<p>I reimplemented Bidaf by myself.<br>
I’ve been having trouble dealing with GPU memory leak</p>
<p>The first image I attached is the codes inside model’s forward function<br>
The table is created by pytorch profiler.</p>
<p>Only a few codes occupy 4GB memory, especially torch.cat.</p>
<p>I’ve been trying to look up the solution to fix it, I haven’t made it.</p>
<p>Is there anyone who knows how to fix it??</p>

ptrblck: <p>Could you explain a bit more why you conclude that <code>torch.cat</code> is leaking memory?<br>
Since <code>torch.cat</code> will create a new tensor, new GPU memory allocation is expected, but it seems as if you are losing memory (leaking) and run eventually into an OOM using <code>torch.cat</code>?</p>

Qmin-Lee: <p>To prevent the model from running into OOM, I set the batch size = 8, instead of using a baseline batch size =60.</p>
<p>Only one torch.cat operation uses 4GB memory looks like there is something wrong.</p>
<p>just so you know : help_h, help_u, temp_s have about (8, 200,100,200) shape</p>

ptrblck: <p>That sounds strange and I cannot reproduce the issue as I get the expected <code>2x</code> increase:</p>
<pre><code class="lang-python">x1 = torch.randn(8, 200, 100, 200, device='cuda')
x2 = torch.rand_like(x1)
x3 = torch.rand_like(x1)

print(torch.cuda.memory_allocated() / 1024**3)
# 0.35762786865234375

y = torch.cat((x1, x2, x3), dim=-1)
print(torch.cuda.memory_allocated() / 1024**3)
# 0.7152557373046875
</code></pre>

Qmin-Lee: <p>I just tried your code and as you said, cat operation doesn’t exceed 1GB.</p>
<p>So, it still remains a mystery. I mean, PyTorch profiler showed me the table and it said 4.74GB is allocated at this operation…</p>
<p>any idea that I should check other things?</p>

ptrblck: <p>I’m not sure if the profiler reporting might be misleading, but are you seeing the same allocations using my code snippet in the profiler?</p>

Qmin-Lee: <p>yes, I confirmed it<br>
the profiler repoting might be wrong, but probably not since I suffer from unexpected GPU memory usage.</p>
<p>According to the original paper and other re-implementation, this model uses under 12gb gpu memory at batch size=60.</p>
<p>But mine uses 14gb at only batch size = 8</p>

Qmin-Lee: <aside class="quote quote-modified" data-post="1" data-topic="8415">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/7ea924/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/torch-cat-blows-up-memory-required/8415">Torch.cat() blows up memory required</a> 
  </div>
  <blockquote>
    Hello! 
When passing the result of a torch.cat() into a Linear layer, the memory required to store the output is much greater than expected. For illustration purposes, I have a simplified example. 
Here I have a matrix where I want to iteratively fill the values after feeding rows of x into a linear layer. I do not use torch.cat() here. The resulting matrix takes about 1 GB memory. 
ff = nn.Linear(100, 1)
x = Variable(torch.randn((60, 200, 100)))
matrix = Variable(torch.zeros(60, 200, 800))

for…
  </blockquote>
</aside>

<p>I’ve read this discussion talking about torch.cat and linear layer(you answered)</p>
<pre><code class="lang-auto">import torch
from torch import nn

x1 = torch.randn(8, 200, 100, 200, device="cuda")

x2 = torch.rand_like(x1)

x3 = torch.rand_like(x1)

print(torch.cuda.memory_allocated() / 1024 ** 3)

# 0.35762786865234375

y = torch.cat((x1, x2, x3), dim=-1)

print(torch.cuda.memory_allocated() / 1024 ** 3)

# 0.7152557373046875

print(y.device)

print(y.shape)

temp = nn.Linear(600, 1)

temp.cuda(0)

print(y.device, x1.device, x2.device, x3.device)

k = temp(y)

print(k.shape)

print(torch.cuda.memory_allocated() / 1024 ** 3)

# 0.73
</code></pre>
<p>I reused your code to check if the linear layer is the cause, but it wasn’t.</p>
<p>I feel like I’ve tried almost everything to solve this problem.  <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>It seems the original issues is:</p>
<aside class="quote no-group" data-username="Qmin-Lee" data-post="7" data-topic="153526">
<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/qmin-lee/48/50344_2.png" class="avatar"> Qmin-Lee:</div>
<blockquote>
<p>According to the original paper and other re-implementation, this model uses under 12gb gpu memory at batch size=60.</p>
<p>But mine uses 14gb at only batch size = 8</p>
</blockquote>
</aside>
<p>If so, could you add the debug print statements to different parts of the code (e.g. after model init, after the foward, after backward, after <code>step()</code> etc.) to see how much memory is used in each step and where the unexpected memory growth might be coming from?</p>