Pytorch crashes on optimizer.step() due to insufficient memory

username: Pravesh_Koirala
cooked: <p>Hello,</p>
<p>I am trying to calculate loss on a function which looks like this:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/f/5f9bb083178fac21277136148b49db60a0a64044.png" alt="image" data-base62-sha1="dDMVdEYS0F5AqJMQieySYi0NlMo" width="176" height="40"></p>
<p>W is a vector of shape (314287, 300), whereas V is a vector of shape (314287, 3000)… (I am trying to learn a sparse representation).</p>
<p>The relevant portion of my code is as follows:</p>
<pre><code class="lang-auto">model = torch.optim.Adam([sparse_embeds], lr=1e-4)
embeds = embeds.cuda()
sparse_embeds = sparse_embeds.cuda()
buckets_size = 1000

def bucketize(tensor, buckets):
  total_buckets = int(tensor.shape[0] / buckets) + 1
  i = 0
  while i*buckets &lt; tensor.shape[0]:
    yield i*buckets, (i+1)*buckets
    i += 1


for e in range(50):
  for i in range(len(embeds)):
    for se in bucketize(embeds, buckets_size):
      start, end = se
      all_dots = embeds[start:end] @ embeds[i]
      s1 = (sparse_embeds[start:end] @ sparse_embeds[i]) - all_dots
      loss = torch.sum((s1.pow(2)))

      print(loss.item())
      model.zero_grad()
      loss.backward()
      model.step()   # &lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt; Memory crashes here.
</code></pre>
<p>I had thought that breaking the loss calculation in multiple buckets would have solved the issue but it doesn’t seem to. Any help at all would be appreciated.</p>

username: albanD
cooked: <p>Hi,</p>
<p>Have you tried to reduce the embedding size or replace Adam with SGD?<br>
The Adam optimizers will allocate a couple buffers with the same size as the weights to store all the statistics it needs. So that might be too much for your machine?</p>

username: Prasanga
cooked: <p>As <a class="mention" href="/u/alband">@albanD</a> suggested try considering SGD instead of Adam as Adam needs more memory for storing first and second moment vectors.</p>
<p>Also, reducing embedding size and bucket_size and overall architecture might help.<br>
Also want to suggest you that, <code>bucketize()</code> only needs tensor_shape instead of the whole tensor.</p>

username: Pravesh_Koirala
cooked: <p>Hey <a class="mention" href="/u/alband">@albanD</a>, I am trying out SGD and it is working. Thank you for the suggestion. A question though. If I am optimizing over a huge tensor with Adam but am only operating on a narrowed view of the tensor at a time. Will the optimization procedure still update (or allocate memory for updating) all the gradients? Can I change this behavior in any way?</p>

username: albanD
cooked: <p>I am not very familiar with what it does in details, but I guess that’s what SparseAdam is for: <a href="https://pytorch.org/docs/stable/optim.html?highlight=adam#torch.optim.SparseAdam">https://pytorch.org/docs/stable/optim.html?highlight=adam#torch.optim.SparseAdam</a> <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: Pravesh_Koirala
cooked: <p>I was not sure if reducing the bucket size helped (I tried with a low value in the first place). Perhaps Adam still updates all tensors (even if they were not involved in gradient computation). Can you suggest me if and how I could change this behavior?</p>

username: Pravesh_Koirala
cooked: <pre><code class="lang-auto">RuntimeError: SparseAdam does not support dense gradients, please consider Adam instead
</code></pre>
<p>I suspect it only works with sparse tensors. I am actually working with dense tensors (to learn a sparse representation). Regardless, I think I will make do with SGD for now. Thank you again <a class="mention" href="/u/alband">@albanD</a></p>

username: Prasanga
cooked: <p>I don’t know the exact problem but optimizer updates only those parameters which leads to loss computation by some means and whose <code>requires_grad</code> is set to True. Try setting  <code>requires_grad</code> of the parameters you don’t want to update, to False or even <code>detach()</code> those parameters from backward computation graph if you also want layers below it  not to be updated.</p>