Retrieve per-class gradients without create_graph=True and retain_graph=True

maljumaily: <p>I am trying to keep track of the per-class gradients of 10k+ classes. Having <code>retain_graph=True</code>allows to access them but uses significant amount of memory. However, I am trying to figure out a way where I could have both <code>create_graph=False</code>and <code>retain_graph=True</code>while having access to per-class gradients. I have attached my sample code below and am able to set <code>create_graph</code>to <code>False</code>. The line to focus on is:<br>
<code>self._scaler.scale(current_loss).backward(create_graph=create_graph, </code><strong><code>retain_graph=True)</code></strong></p>
<p>Any help is appreciated. Thank you in advance!</p>
<pre data-code-wrap="python"><code class="lang-python">

class NativeGrad:
    state_dict_key = "amp_scaler"

    def __init__(self):
        self._scaler = torch.amp.GradScaler('cuda')

    def __call__(self, loss, optimizer, clip_grad=None, model=None, create_graph=False, update_grad=True):
        grads: list = []
        for current_loss in loss:
            for p in model.parameters():
                if p.grad is not None:
                    p.grad.data.zero_()

            self._scaler.scale(current_loss).backward(create_graph=create_graph, retain_graph=True) # ← ← ← retain_graph should be False here

            grad = torch.cat([
                p.grad.flatten().clone() if p.grad is not None else torch.zeros_like(p).flatten()
                for p in model.parameters()
            ])
            grads.append(grad)

        for p in model.parameters():
            if p.grad is not None:
                p.grad.data.zero_()

        grads_tensor = torch.stack(grads, dim=0)
        
        # Manipulate the grads_tensor...
        

    def state_dict(self):
        return self._scaler.state_dict()

    def load_state_dict(self, state_dict):
        self._scaler.load_state_dict(state_dict)


</code></pre>

valerian.rey: <blockquote>
<p>Having <code>retain_graph=True</code> allows to access them but uses significant amount of memory.</p>
</blockquote>
<p>If you’re concerned about memory usage, be aware that 1 gradient takes as much memory as the entire model, so storing 10k gradients is gonna cost you an amount of memory that is generally prohibitive (no matter how you compute them).</p>
<p>In your code example, you compute these gradients sequentially, in a for loop. If you do that, you have to retain the graph because subsequent gradient computations will still need access to the computation graph. But this is not what makes your memory usage blow up: you just keep the computation graph that was built during the forward pass (and that is not that big at all). Most of the memory usage (I guess 99.9%+ if you have 10k+ gradients) should come from your <code>grads</code> variable.<br>
So I think there’s no direct solution to your problem.</p>
<p><code>create_graph</code>, on the other hand, would make your memory usage blow up even harder, because for each of those 10k+ gradients you would have to store the computation graph leading to them. But you should only use <code>create_graph</code> if you care about higher order differentiation (which is not gonna help you because it’s memory expensive and you already struggle with memory).</p>

maljumaily: <p>Thank you, Valerian, for your response. My entire objective is getting the individual gradient of each class. Given the situation, are there any optimization technique (or approach) that could be used to make the computation slower but minimizes the usage of the GPU memory? Even approximations of the gradients could be viable. I look forward to hearing your feedback!</p>

valerian.rey: <p>It depends on what you want to do with your gradients.</p>
<p>I have a similar setup in my research project where I’m actually interested in the inner products between all pairs of individual gradients (be it per-class, per-sample, per-loss or per-whatever gradients). So in my case, I can just compute these inner products as sum of partial inner products and I don’t ever have to store the full gradients in memory (see algo 3 of <a href="https://arxiv.org/pdf/2406.16232" rel="noopener nofollow ugc">our article</a> if this seems related to your use-case - although I don’t think it can be efficiently implemented for per-class gradients). If you’re rather interested in the norms of these gradients, a similar approach exists (e.g. <a href="https://arxiv.org/pdf/1510.01799" rel="noopener nofollow ugc">https://arxiv.org/pdf/1510.01799</a>). In general, I would say that it really depends on what statistic you want to compute from these gradients. And depending on this statistic, an optimized algorithm may exist to compute it without ever storing the full gradients in memory. There exists a library called <a href="https://docs.backpack.pt/en/master/extensions.html#first-order-extensions" rel="noopener nofollow ugc">backpack</a> that has algorithms to compute a few other statistics that you could compute from those individual gradients (among other things). I’m not sure it would solve your problem because this library only works for a limited subset of pytorch layers, and it might not work for per-class gradients. But it could maybe give you some ideas.</p>
<p>A few ideas that come to my mind if you really need the raw gradients:</p>
<ul>
<li>Reduce floating point precision / quantize the gradients. This could reduce memory usage by maybe a factor 2 or 4.</li>
<li>Group some classes together: if the gradients of two classes are generally very aligned (cosine similarity close to 1), you could probably group those two classes together.</li>
<li>Only compute the gradients with respect to a subset of the parameters of the model (e.g. the last layer’s parameters).</li>
<li>Once enough gradients have been computed in your for-loop to fill up the GPU memory, store them on disk, and stop storing them on GPU memory. Repeat that until you have stored all the gradients on disk. This will be slow, so you shouldn’t rely on that if performance is important to you.</li>
</ul>
<p>I hope this helps! Also, please share more details about your goals if you can. I think this is an interesting discussion <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=14" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

maljumaily: <p>Thank you so much for your response. Surprisingly, your paper is one of the papers my supervisor assigned for me to go over. I have been reading it and the premise of what I want is similar. I have to get the per-class gradient and update them accordingly (we haven’t talked in-depth about the main goal here).</p>
<p>What threw me off now is the fact that storing one gradient is the equivalent of storing all gradients. Am I able to store the gradients of the current minibatch only? If not, then I will have to go into the recommendations you have suggested.</p>

valerian.rey: <blockquote>
<p>Surprisingly, your paper is one of the papers my supervisor assigned for me to go over.</p>
</blockquote>
<p>Oh that’s funny!</p>
<blockquote>
<p>What threw me off now is the fact that storing one gradient is the equivalent of storing all gradients.</p>
</blockquote>
<p>Well, in our case (and perhaps in your case as well), we end up doing a linear combination of the gradients. So if you know the coefficients with which you wanna make this linear combination (one coefficient per gradient), you have two possibilities:</p>
<ol>
<li>Compute the linear combination of the gradients by multiplying each gradient by the corresponding coefficient and summing the results. If you store all gradients in memory at some point, this takes a ton of memory.</li>
<li>Compute the linear combination of the different elements of the loss vector scaled by the corresponding coefficient, and then compute the gradient of the obtained loss. By linearity of differentiation, the result will be the same as for approach 1), but you can save a ton of memory by doing that.</li>
</ol>
<p>I’m really not sure you can use the same method as we have - this will really depend on how your coefficients are computed, if you even have coefficients to begin with.</p>
<blockquote>
<p>Am I able to store the gradients of the current minibatch only?</p>
</blockquote>
<p>I’m not sure to understand this question. Do you currently store gradients accross mini-batches? This seems like an even greater memory requirement than just using the information from the current mini-batch to update the parameters.</p>