Exploding memory

Surya_Narayanan: <p>I am facing an issue where my memory usage is exploding, and I can’t explain why. Here is what my code looks like</p>
<pre><code class="lang-auto">modes = torch.randn(num_models, 1, 512, 30522).to(device)

for epoch in range(num_epochs):
    model.train()

    for batch_ix, batch in enumerate(train_dataloader):
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch) 
        
        distances = torch.cat([((outputs.logits - mode)**2).sum(axis=[1,2]) for mode in modes])

        chosen_model_ix = torch.argmin(distances)

        #update the new mode
        modes[chosen_model_ix] += outputs.logits
        modes[chosen_model_ix] /= 2

</code></pre>
<p>After playing with <code>torch.no_grad()</code>, it seems these lines are making the most difference to the GPU memory footprint</p>
<pre><code class="lang-auto">#update the new mode
modes[chosen_model_ix] += outputs.logits
modes[chosen_model_ix] /= 2
</code></pre>
<p>Does this make sense? Why would these inplace operations blow up my memory?</p>

ptrblck: <p>It seems you are accumulating tensors which are still attached to a computation graph and are thus disallowing PyTorch to free any intermediate tensors.<br>
I don’t know if <code>modes</code> should stay differentiable, but in case that’s not the case and it’s used as a static target/reference, you should <code>.detach()</code> the logits before accumulating them.</p>

Surya_Narayanan: <p>Thanks, that seemed to curb the memory consumption. I"ll follow up with further queries if that didn’t satisfy my research problem.</p>

Surya_Narayanan: <p>Hey, just to reopen this issue, is there any reason that calculating accuracy should cause GPU memory usage to blow up?</p>
<p>This is the relevant part of my code- I’ve narrowed the bug down to the line that calculates accuracy- this seems to be the part of the code that’s blowing up the GPU memory.</p>
<pre><code class="lang-auto">model_cx_accuracy = Accuracy(task="multiclass", num_classes = len(models)).to(device)
mod_pred_class = model_ops.logits.argmax(dim = -1)
mod_mo_acc = gen_accuracy(mod_pred_class, batch['labels']) # Buggy line
</code></pre>
<p>Does this make sense?</p>

ptrblck: <p>It depends what exactly <code>gen_accuracy</code> does so did you check which tensors are allocated inside this function and how large these are?</p>

Surya_Narayanan: <p>whoops, I intended to paste this earlier but forgot-</p>
<pre><code class="lang-auto">gen_accuracy = Accuracy(task="multiclass", num_classes=30522, ignore_index=-100, average = 'micro', multidim_average = 'samplewise').to(device)
</code></pre>
<p>the tensors are 8 * 512 * 30522 each, but i don’t see why that space isn’t being allocated and freed each time <code>gen_accuracy</code> is being called</p>

ptrblck: <p>The shape wouldn’t make sense for a<code>multiclass</code> classification as the class dimension is expected to be in <code>dim1</code>. In any case, using your shapes the method seems to use ~500MB:</p>
<pre><code class="lang-python">import torch
from torchmetrics import Accuracy

print(torch.cuda.memory_allocated()/1024**2)
# 0.0

device = "cuda"
gen_accuracy = Accuracy(task="multiclass", num_classes=30522, ignore_index=-100, average = 'micro', multidim_average = 'samplewise').to(device)

x = torch.randn(8, 30522, 512, device=device)
y = torch.randint(0, 30522, (8, 512), device=device)
print(torch.cuda.memory_allocated()/1024**2)
# 476.9375

loss = gen_accuracy(x, y)
print(torch.cuda.memory_allocated()/1024**2)
# 484.39111328125
</code></pre>

Surya_Narayanan: <p>Yes, I think I <code>argmax</code> over the last dimension and then feed in two <code>8 x 512</code> vectors.</p>

Surya_Narayanan: <p>But, this doesn’t explain why the accuracy is exploding, right? It seems this accuracy costs 9 MB or so MB every time it’s called. I think that’s what’s costing me to run out of memory.</p>
<pre><code class="lang-auto">for i in range(100):
  loss = gen_accuracy(preds, y)
  print(torch.cuda.memory_allocated()/1024**2)
</code></pre>
<p>Is there anything I can do to prevent my accuracy function hogging a new memory chunk memory each time it’s called?</p>

Surya_Narayanan: <p>it’s interesting that using the Accuracy() function blows up memory, but manually implementing it, doesn’t.</p>
<pre><code class="lang-auto">import torch
from torchmetrics import Accuracy

print(torch.cuda.memory_allocated()/1024**2)
# 0.0

device = "cuda"


x = torch.randn(8, 512, 30522, device=device)
preds = x.argmax(dim = -1)
y = torch.randint(0, 30522, (8, 512), device=device)
print(torch.cuda.memory_allocated()/1024**2)
# 476.96875

for i in range(100):
  gen_accuracy = Accuracy(task="multiclass", num_classes=30522, ignore_index=-100, average = 'micro', multidim_average = 'samplewise').to(device)
  acc1 = torch.sum((preds == y).float(), dim = -1) #doesn't blow up memory
  acc2 = gen_accuracy(preds, y) #does blow up memory
  print(torch.cuda.memory_allocated()/1024**2)

</code></pre>
<p>Any thoughts on whether I’m missing something?</p>

ptrblck: <p>I don’t think you are missing anything and I can also reproduce the increased memory usage.<br>
However, I needed to jump on another issue and didn’t debug the issue yet.<br>
So far I would blame the <code>Accuracy</code> method and would use a manual implementation.</p>

Surya_Narayanan: <p>ok, let me also report this on torchmetrics’ github, see if anyone has any thoughts</p>