Pytorch appears to be crashing due to OOM prematurely?

sad_robot: <p>I’ve got a model with 110M parameters, and I’m training it on a very small dataset (like 500 examples).</p>
<p>Yet that is enough to crash Pytorch on a K80 GPU with 11GB VRAM.<br>
What is going on here? 110M x 4 (float size) = 440M = 0.440 GB + minuscule dataset size != 11GB VRAM…</p>
<p>Thanks for your help</p>

gphilip: <aside class="quote no-group" data-username="sad_robot" data-post="1" data-topic="131039">
<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/sad_robot/48/25510_2.png" class="avatar"> sad_robot:</div>
<blockquote>
<p>110M x 4 (float size) = 440M = 0.440 GB</p>
</blockquote>
</aside>
<p>I would be very surprised if just the float size is taken up by your model. Especially when it is training, it has to keep track of gradients. And there are probably lots of other things which need to be in memory for the model to work.</p>
<p>See if you get the out of memory error if you train the network on a batch with just one input. If that works, try increasing the number of inputs in a batch till you figure out the breaking point.</p>
<p>Another thing: perhaps the GPU memory is being used by other processes, so that not all of the 11GB is available? The output of <code>nvidia-smi</code> should tell you if this is the case or not.</p>

sad_robot: <aside class="quote no-group" data-username="gphilip" data-post="2" data-topic="131039">
<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/g/858c86/48.png" class="avatar"> gphilip:</div>
<blockquote>
<p>should</p>
</blockquote>
</aside>
<p>You think those other things could increase the memory burden roughly x20?</p>
<p>Regarding your Suggestions:<br>
I have tried increasing the batch size slowly, it works with very small batch sizes like 10 I think.<br>
And yes I’ve checked with nvidia-smi, I’m the only one using this machine and there are no zombie processes.</p>
<p>Extra Context:<br>
Also I noticed there could be a memory leak in my code, in the sense that the mini batch tensor is reallocated at each update step &amp; it may not be garbage collected quickly enough. Still I tried fixing this with a reusable mini-batch Tensor and it still didn’t work…</p>

gphilip: <p>How big is each element of the input data set? Perhaps it is the data loader which is taking up lots of memory?</p>

sad_robot: <p>Each element is 200 floats so 800 bytes.<br>
Here is the brief version of my code (with only minimal memory leak patches):</p>
<pre><code class="lang-auto"># pad_len := max number of tokens (input sequences padded to this length)
def train_loop(model, input_output_data, epochs=5, batch_size=32, pad_len=200):
    input, output = input_output_data
    n_examples = len(input)
    
    model.train() # turn on training
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    n_batches = int(n_examples/batch_size+0.99999)
    for i in range(epochs):
        print(f'epoch: {i+1}/{epochs}')
        for j in range(n_batches):
            optimizer.zero_grad()
            torch.cuda.empty_cache()
            try:
                # they must be padded to the same size for batching to work...
                inputs = tokenizer(input[j*batch_size:(j+1)*batch_size], return_tensors='pt',
                                   padding='max_length', truncation=True, max_length=pad_len)
                outputs = tokenizer(output[j*batch_size:(j+1)*batch_size], return_tensors='pt',
                                    padding='max_length', truncation=True, max_length=pad_len)
                loss = model(**inputs, labels=outputs['input_ids'])[0].mean()
            except Exception as e:
                #pdb.set_trace()
                raise
            
            print(f'batch: {j+1}/{n_batches}, loss: {loss}')
            loss.backward() # mean is unnecessary but just for safety
            optimizer.step()
</code></pre>

gphilip: <p>I can’t see anything in this code that suggests large memory use. As a band-aid I can suggest that you could try adding</p>
<pre><code class="lang-auto">del inputs
del outputs
del loss
</code></pre>
<p>after the line</p>
<pre><code class="lang-auto">optimizer.step()
</code></pre>
<p>.</p>

ptrblck: <aside class="quote no-group" data-username="sad_robot" data-post="3" data-topic="131039">
<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/sad_robot/48/25510_2.png" class="avatar"> sad_robot:</div>
<blockquote>
<p>You think those other things could increase the memory burden roughly x20?</p>
</blockquote>
</aside>
<p>The memory usage is model-dependent and often the majority of the memory is used by the forward activations, not the parameters or gradients.<br>
E.g. in <a href="https://discuss.pytorch.org/t/how-to-deal-with-excessive-memory-usages-of-pytorch/126098/5">this post</a> I’ve posted some stats about the used model and you can see that the activations use:</p>
<pre data-code-wrap="python"><code class="lang-python">6452341200 / 138357544 ~= 47
</code></pre>
<p>times more memory than the parameters.</p>

sad_robot: <aside class="quote no-group" data-username="ptrblck" data-post="7" data-topic="131039">
<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>E.g. in <a href="https://discuss.pytorch.org/t/how-to-deal-with-excessive-memory-usages-of-pytorch/126098/5">this post</a> I’ve posted some stats about the used model and you can see that the activations use:</p>
<pre><code class="lang-auto">6452341200 / 138357544 ~= 47
</code></pre>
<p>times more memory than the parameters.</p>
</blockquote>
</aside>
<p>Wow that’s astonishing. I heard it was more but I had no idea it was this much more.</p>

ptrblck: <p>It really depends on the model architecture and especially for e.g. conv layers, you would see a huge memory difference, while linear layers could yield the inverse effect.<br>
Here is a smaller example:</p>
<pre><code class="lang-python"># conv
model = nn.Conv2d(3, 64, 3, 1, 1)
x = torch.randn(1, 3, 224, 224)

out = model(x)

model_param_size = sum([p.nelement() for p in model.parameters()])
input_size = x.nelement()
act_size = out.nelement()

print('model size: {}\ninput size: {}\nactivation size: {}'.format(
    model_param_size, input_size, act_size))

&gt; model size: 1792
  input size: 150528
  activation size: 3211264
  
# linear
model = nn.Linear(1024, 1024)
x = torch.randn(1, 1024)

out = model(x)

model_param_size = sum([p.nelement() for p in model.parameters()])
input_size = x.nelement()
act_size = out.nelement()

print('model size: {}\ninput size: {}\nactivation size: {}'.format(
    model_param_size, input_size, act_size))

&gt; model size: 1049600
  input size: 1024
  activation size: 1024
</code></pre>

sad_robot: <aside class="quote no-group" data-username="ptrblck" data-post="9" data-topic="131039">
<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>e.g. conv layers, you would see a huge memory difference, while linear layers could yield the inverse effect.</p>
</blockquote>
</aside>
<p>Hmm that’s a good point because of weight sharing… Does my transformer have weight sharing (I’m using <a href="https://arxiv.org/abs/1906.08237" rel="noopener nofollow ugc">XLNet</a>)? I think yes because of recurrence right?</p>
<p>Also I guess the remedy is pipeline parallelism right?</p>

ndvbd: <p>@ <a href="https://discuss.pytorch.org/u/ptrblck">ptrblck</a> is there a better way than out.nelement() to find the real activation size of a network? After all, consider for example a UNET, the out.nelement()  only shows the number of elements in the output tensor and not in all tensors along the way. I want a way to estimate how increasing the batch size increases the memory consumption, and would like to know the real size of activations during forward pass.</p>

ptrblck: <p><code>nelement()</code> returns the number of elements of a tensor. If you want to check intermediate forward activation sizes you could check them separately as described e.g. in <a href="https://discuss.pytorch.org/t/resnet-50-takes-10-13gb-to-run-with-batch-size-of-96/117402/2">this post</a>.</p>

ndvbd: <p>Thanks. I wrote this function to calculate the activation size of a network:</p>
<pre><code class="lang-auto">
total_output_elements = 0
def calc_total_activation_size(model, call_the_network_function):
    global total_output_elements
    total_output_elements = 0

    def hook(module, input, output):
        global total_output_elements
        total_output_elements += output.numel()
        
    handle = torch.nn.modules.module.register_module_forward_hook(hook)
    result = call_the_network_function()
    handle.remove()
    return result, total_output_elements
</code></pre>