Bizzare PyTorch CUDA memory allocation failure on Linux

calebh: <p>I am encountering a bizarre CUDA memory allocation error on Linux (and not Windows). I have two different machines - a Windows machine with an NVIDIA GeForce GTX 1050 with 4 GB of RAM and a Google Cloud NVIDIA Tesla T4 with 16 GB of RAM. My model (linked below) is a transformer which have a reputation for chewing up a lot of memory (see the Reformer paper for example).</p>
<p>On my Windows machine on the latest PyTorch version I can run the model fine in GPU memory for most training examples with a batch size of 4 (the training examples are variable sized). When I take the exact same model and training data over to the Linux VM I begin getting CUDA out of memory errors after around 40 steps. Increasing the batch size decreases the number of steps needed to obtain an out of memory error. This would indicate a memory leak to me, so I went ahead and made sure I wasn’t holding onto any references to allow the garbage collector to do its thing. I have added calls for clearing the GPU cache, running the Python garbage collector, etc. which doesn’t seem to help.</p>
<p>The super weird part is if I insert a call to</p>
<pre><code class="lang-auto">print(torch.cuda.memory_summary(device=None, abbreviated=True))
</code></pre>
<p>inside my exception handler for the CUDA out of memory error the issue goes away! I am even able to bump the batch size to 12 on the Linux instance, which is what I would expect from a system with higher GPU memory.</p>
<p>Model is here:<br>
<aside class="onebox githubgist">
  <header class="source">
      <a href="https://gist.github.com/calebh/fd7ff5014d9925b0e2d7473d9d4d5131#file-model-py-L257" target="_blank" rel="nofollow noopener">gist.github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://gist.github.com/calebh/fd7ff5014d9925b0e2d7473d9d4d5131#file-model-py-L257" target="_blank" rel="nofollow noopener">https://gist.github.com/calebh/fd7ff5014d9925b0e2d7473d9d4d5131#file-model-py-L257</a></h4>
<h5>model.py</h5>
<pre><code class="Python">import torch
import pickle
import os
import random
import sourcenode
import torch.nn as nn
import torch.utils.checkpoint
import torch.utils.data
import math
import objprocessor</code></pre>
This file has been truncated. <a href="https://gist.github.com/calebh/fd7ff5014d9925b0e2d7473d9d4d5131#file-model-py-L257" target="_blank" rel="nofollow noopener">show original</a>

<p>
</p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Any ideas?</p>

calebh: <p>Update:</p>
<p>If I stick</p>
<pre><code class="lang-auto">torch.cuda.memory_summary(device=None, abbreviated=True)
</code></pre>
<p>in a loop that I iterate over 25 times I can push up to a batch size of 16. Sleeping and waiting in a busy loop do not reproduce the same behaviour.</p>
<p>Without looping over this function invocation multiple times I see periods where I will get many CUDA allocation failures in a row. Looping results in occasional CUDA allocation failures which is what I would expect due to sometimes encountering input sequences of high length.</p>

ptrblck: <aside class="quote no-group" data-username="calebh" data-post="2" data-topic="79571">
<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/calebh/48/22756_2.png" class="avatar"> calebh:</div>
<blockquote>
<p>Sleeping and waiting in a busy loop do not reproduce the same behaviour.</p>
</blockquote>
</aside>
<p>Could you explain a bit, what that means?<br>
Does it mean, that running <code>torch.cuda.memory_summary</code> allows you to increase the batch size, while a sleep call does not?</p>

calebh: <p>That is correct. I have no idea why a call to <code>torch.cuda.memory_summary</code> stops the out of memory errors, but I am glad it does. Basically when I see a CUDA out of memory error I call <code>memory_summary</code> 25 times in a row then try one more time to feed the example to the network. Most of the time this fixes the out of memory error.</p>

ptrblck: <p>That’s indeed weird, as it should only call to <a href="https://github.com/pytorch/pytorch/blob/1f09f7ea448f404fff135b5e4c1c3d63d157de79/torch/csrc/cuda/Module.cpp#L286">this method</a> to get the stats.<br>
Does a <code>torch.cuda.synchronize()</code> instead of <code>torch.cuda.memory_summary()</code> yield the same behavior?</p>