Modify the autograd backend to propogate gradients to only one child during backwards pass

username: Aradhye_Agarwal
cooked: <p>I am working on PEFT techniques of LLMs. An idea which I wanted to try out was to try out full finetuning while only propogating the maximum gradient from any particular node. That is, given a node N in the computation graph on which N.backward() has been called, first compute the gradients for the children of N, and then only pass the gradient and call backward() on the child having the max gradient (take L1 norm in case of non-scalar parameters). This would significantly reduce the backward() pass overhead.</p>
<p>However, this is proving to be very difficult because pytorch’s autograd uses the C++ engine as its backend and that too its compiled binaries. Also it is tricky to figure out where to make the change (I tried looking into the codebase but it seems that the backend does not do the backward pass recursively but rather iteratively using a multitude of datastructures), so there is no clean way of not calling backward on the nodes having non max gradient.</p>
<p>Some guidance on how to achieve the desired functionality would be helpful. Please note that speed here is not an issue, I just want to assess how the model trained (or finetuned) using such a technique would perform. Hence it would be fine if the backend used is a python engine. I however, can’t seem to figure out how to do this.</p>
<p>Thanks</p>

username: KFrank
cooked: <p>Hi Aradhye!</p>
<aside class="quote no-group" data-username="Aradhye_Agarwal" data-post="1" data-topic="193818" data-full="true">
<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/aradhye_agarwal/48/66299_2.png" class="avatar"> Aradhye_Agarwal:</div>
<blockquote>
<p>first compute the gradients for the children of N, and then only pass the gradient and call backward() on the child having the max gradient</p>
</blockquote>
</aside>
<p>Depending upon the details of what you want to do, you might be able to<br>
use <a href="https://pytorch.org/docs/2.1/generated/torch.autograd.grad.html" rel="noopener nofollow ugc">autograd.grad()</a>.  Note, this is part of pytorch’s python interface, so<br>
there would be no need to drill down into the C++ implementation.</p>
<p>However, it’s not clear to me exactly what you want.  Let’s say that the last<br>
layer of your model is a <code>Linear</code> and you pass its output to some loss function,<br>
and you imagine calling <code>loss.backward()</code>.  <code>loss</code> will typically have one child,<br>
namely the output of your final <code>Linear</code>.  Conceptually, this will have three<br>
children – the <code>Linear</code>’s <code>weight</code>, its <code>bias</code>, and the input to that <code>Linear</code>, likely<br>
the output of the previous layer.</p>
<p>What specifically do you propose to do?  The <code>weight</code> and <code>bias</code> have no<br>
children, but are trainable, while the input has children – the parameters<br>
of the previous layer – but is not, itself, trainable.</p>
<p>Let’s say that of these three children, <code>weight</code> turns out to have the largest<br>
gradient.  Do you now populate <code>weight.grad</code>, but otherwise backpropagate<br>
no further?  On the other hand, let’s say that the <code>Linear</code>’s input has the<br>
largest gradient.  Do you then continue your “max-gradient” backpropagation<br>
through the previous layer, but not accumulate a gradient into <code>weight.grad</code><br>
nor fine-tune <code>weight</code>?</p>
<p>Also let’s say that if you were to perform the full backpropagation, you would<br>
discover, for example, that the <code>bias</code> of the first layer has a gradient larger<br>
than the <code>weight</code> of the last layer.  If you stopped backpropagation at the last<br>
layer as soon as you determined that the last layer’s <code>weight</code> had the largest<br>
gradient of the three children in question, you would never discover that the<br>
first layer’s weight had a larger gradient.  Would such a scenario fit properly<br>
into your use case?</p>
<p>Best.</p>
<p>K. Frank</p>

username: Aradhye_Agarwal
cooked: <p>Hi Frank,</p>
<p>Actually what I wanted was this:</p>
<p>During back propogation, when finding derivative of loss w.r.t the current parameter, only take the scalar gradients (where scalars are the individual elements in the tensor parameter) having magnitude in top k of the scalar gradients. Fix the rest to 0.</p>
<p>Note, that doing this during optimizer.step() would not be correct, a naive way of setting the param.grad to 0 for all but the top k magnitude gradients in the current param would not produce the required effect since the all the params in the succeeding layers have been used for gradient computation, when only the top k gradients were supposed to be used.</p>
<p>As a result I would need to modify the backward() pass to set the non topk gradients to 0 at each node. Also note that one motivation for this is increase the speed of the backward pass. And hence it would make sense to pass sparse gradients (since most would be set to 0) to the children which would then do the same thing and so on.</p>
<p>I am not quite able to figure out how to achieve this. Since speed is crucial, I may also have to modify the C++ engine. I need to do this to compare the speed increase compared to the “dense” gradient computation (ie. the default autograd behaviour). Note that this is different from my previous question where speed was not an issue.</p>
<p>Looking forward to your reply.</p>

username: KFrank
cooked: <p>Hi Aradhye!</p>
<aside class="quote no-group" data-username="Aradhye_Agarwal" data-post="3" data-topic="193818" data-full="true">
<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/aradhye_agarwal/48/66299_2.png" class="avatar"> Aradhye_Agarwal:</div>
<blockquote>
<p>I would need to modify the backward() pass to set the non topk gradients to 0 at each node.</p>
</blockquote>
</aside>
<p>You should be able to do this by registering a <a href="https://pytorch.org/docs/2.1/generated/torch.Tensor.register_hook.html" rel="noopener nofollow ugc">hook</a> on each node to which<br>
you want to pass zeroed-out gradients.  The hook would perform the topk<br>
computation and create a new gradient tensor with the non-topk entries set<br>
to zero.</p>
<aside class="quote no-group">
<blockquote>
<p>Also note that one motivation for this is increase the speed of the backward pass. And hence it would make sense to pass sparse gradients</p>
</blockquote>
</aside>
<p>I don’t whether there is a clean way to do this.  If you pass a gradient that is<br>
packaged as a full tensor into a backward hook and have hook return a<br>
sparse-tensor gradient, autograd will complain.  I don’t know whether you can<br>
arrange things so that nothing but sparse gradients flow through you hooks.</p>
<p>You can pass a sparse tensor into <code>autograd.grad()</code> as its <code>grad_outputs</code><br>
argument (the vector in the vector-jacobian product that <code>autograd.grad()</code><br>
computes), so perhaps if you perform the backward pass “manually” by<br>
chaining together a bunch of <code>autograd.grad()</code> calls, you could work with<br>
sparse gradients.</p>
<p>Having said that, working with sparse gradients may or may not give you<br>
any speed-up (and could even slow things down).  The reason is that even<br>
though your sparse-gradient computations will use (potentially significantly)<br>
fewer floating-point operations, the gpu (and even cpu) floating-point pipelines<br>
are so heavily optimized for structured floating-point arithmetic (e.g., matrix<br>
multiplication) that the speed they lose on the less-structured sparse-tensor<br>
operations could more than undo any speed benefit of doing “less” work.  The<br>
only good way I know to evaluate these trade-offs is to time concrete code<br>
(and those relative timings may differ significantly across hardware platforms).</p>
<p>Good luck!</p>
<p>K. Frank</p>