Can you override torch.ops.aten._backward using __torch_dispatch__?

kdu: <p>Hi, I’m looking to override the gradient computation (ideally in Python) at the level of the computation graph (e.g. not for a specific op like add or sin, but in the actual backpropagation process through the graph, especially: (a) when the derivative is computed for each node along a path in the computation graph via multiplication and (b) when the derivative is accumulated at a node from multiple paths by addition).</p>
<p>I have been playing around wtih <a href="https://github.com/albanD/subclass_zoo/blob/main/inner_autograd_tensor.py" rel="noopener nofollow ugc">https://github.com/albanD/subclass_zoo/blob/main/inner_autograd_tensor.py</a> to try to  override different parts of the autograd engine.</p>
<p>But, while I can easily override individual operations like embedding or add, it appears that <code>torch.ops.aten._backward</code> is not registered as a func during the call to  <strong>torch_dispatch</strong> when a user calls .backward() on a tensor in Pytorch.</p>
<p>I’ve included a short excerpt of the code I’m playing around with below. This part specifically contains the implementation of <code>__torch_dispatch__</code> for a wrapper around torch.Tensor which can override individual ops but not the backward call.</p>
<p>For the full, runnable example, see <a href="https://colab.research.google.com/drive/11a9uW3b-30j9Fgs5L0NAkKnpvrx-fGuN?usp=sharing" rel="noopener nofollow ugc">this colab notebook</a>.</p>
<pre><code class="lang-auto">@classmethod
    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
        def unwrap(t):
            if isinstance(t, cls):
                return t.elem
            elif isinstance(t, torch.Tensor) and t.requires_grad:
                # If any other argument at this level does require gradients
                # it will not interact with our inner Tensor and thus this
                # should fail.
                raise RuntimeError("Bad mixup of autograd level")
            else:
                return t

        def wrap(t):
            # Micro-optimization: not necessary to rewrap if the output tensor
            # doesn't require gradients
            if (
                isinstance(t, torch.Tensor)
                and not isinstance(t, cls)
                and t.requires_grad
            ):
                return cls(t)
            else:
                return t

        with enable_reentrant_dispatch():
            # Override gradient behavior
            if func == torch.ops.aten.embedding.default:
                args = fill_defaults(args, 5, [-1, False, False])
                weight, indices, padding_idx, scale_grad_by_freq, _sparse = map(
                    unwrap, args
                )
                assert not kwargs
                # Force sparse gradients.  We could have also done this by
                # defining a custom autograd function.
                return cls(func(weight, indices, padding_idx, scale_grad_by_freq, True))
            if func == torch.ops.aten.add.Tensor:
                # This operation is successfully overridden
                print("OVERRIDING ADD")
                t1, t2 = map(
                    unwrap, args
                )

                assert not kwargs
                def override_add(t1, t2):
                  return func(t1, t2) / 10
                return cls(override_add(t1, t2))
            
            if func == torch.ops.aten._backward.default:
                # this never gets called!
                print("OVERRIDING BACKWARD") 

            return tree_map(
                wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs))
            )
</code></pre>
<p>How can I (if at all possible) implement <code>__torch_dispatch__</code> so that the backward call is overridden and controllable by the user?</p>
<p>Thanks for the help!</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/kdu">@kdu</a>,</p>
<p>Have you had a look at using hooks for solving your problem? That can allow for in-place modification of gradients within the Python API. The docs are here: <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_hook" class="inline-onebox" rel="noopener nofollow ugc">Module — PyTorch 2.0 documentation</a></p>

kdu: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>, thanks for the reply! I did take a look and I’m not sure if it fits my use case. To my understanding, it looks like the <code>register_full_backward_pre_hook</code> allows the user to define a function executed before backprop is run through the graph of the module and <code>register_full_backward_hook</code> applies a function after backpropping through the graph of the module. While these hooks are useful to modifying the gradient of the figurative “start” and “end” of the module, it doesn’t seem to let me modify the gradients of nodes in between.</p>
<p>That is, I’d like to be able to modify how the backprop actually executes through each node within the module, but I think these hooks only give me access to modifying the gradient of the head node(s) of the module (<code>grad_output</code> of <code>register_full_backward_pre_hook</code>) and the gradient of the input nodes (<code>grad_input</code> of <code>register_full_backward_hook</code>). Does that match your understanding of what hooks can/cannot do?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/kdu">@kdu</a>,</p>
<aside class="quote no-group" data-username="kdu" data-post="3" data-topic="185889">
<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/kdu/48/62978_2.png" class="avatar"> kdu:</div>
<blockquote>
<p>Does that match your understanding of what hooks can/cannot do?</p>
</blockquote>
</aside>
<p>The hooks allow you to access the gradient of the loss with respect to the output (i.e. <code>grad_output</code>), or the inputs to a module (<code>input</code>) for the <code>full_backward</code> and <code>forward_pre</code> hooks respectively.</p>
<aside class="quote no-group" data-username="kdu" data-post="3" data-topic="185889">
<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/kdu/48/62978_2.png" class="avatar"> kdu:</div>
<blockquote>
<p>gradients of nodes</p>
</blockquote>
</aside>
<p>I assume when you state nodes here, it’s nodes within the graph itself and not nodes within the network? If so, you can still use hooks, but you’ll need to use the <code>torch.autograd.grad.Node.register_hook</code> function, I think. The docs are here: <a href="https://pytorch.org/docs/stable/generated/torch.autograd.graph.Node.register_hook.html?highlight=hook#torch.autograd.graph.Node.register_hook" class="inline-onebox" rel="noopener nofollow ugc">torch.autograd.graph.Node.register_hook — PyTorch 2.1 documentation</a></p>

kdu: <p>Thank <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> - I’ve been playing around with it and I think it might be the answer, but I’m not sure if it lets me a few things, in particular whether I can override how the gradient of a node is accumulated from the gradients of the parent node. I opened a separate question here about that because it felt sufficiently self-contained <a href="https://discuss.pytorch.org/t/access-gradient-of-parents-of-a-node-during-first-execution-of-backprop/186090" class="inline-onebox">Access gradient of parents of a node during first execution of backprop</a>.</p>