Can the PyTorch autograd engine track and correctly compute gradients for a composite custom autograd function?

cannoli: <p>I want to define a custom autograd function that combines multiple contiguous layers into a single staged operation to reduce kernel launches from an external library being used to implement my custom forward/backward pass logic. Specifically, I’m looking to combine several custom autograd functions into one.</p>
<p>My question is: Can the PyTorch autograd engine track and correctly compute gradients for such a composite custom autograd function? I had done a basic implementation where I have a wrapper class that inherits from <code>nn.Module</code> and its forward method accumulates the function calls (where the functions are instances of classes inheritng from <code>torch.autograd.Function</code>). However, when I looked at the computational graph, the custom backward operators I’ve implemented never get called</p>

ptrblck: <aside class="quote no-group" data-username="cannoli" data-post="1" data-topic="206286">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/97f17d/48.png" class="avatar"> cannoli:</div>
<blockquote>
<p>where the functions are instances of classes inheritng from <code>torch.autograd.Function</code>). However, when I looked at the computational graph, the custom backward operators I’ve implemented never get called</p>
</blockquote>
</aside>
<p>This is weird since your custom <code>autograd.Function</code>s should be used including the <code>backward</code> method. Could you post the model definition reproducing this issue, please?</p>

cannoli: <p>Thank you for offering to help! Here is a general idea.</p>
<pre><code class="lang-auto">class ExternalLibraryFunction(Function):
    """
    Base class For User Defined Layers using External Library API
    """

    #The user implements these methods 
    @staticmethod
    def forward(*inputs, **kwargs): ...

    @staticmethod
    def setup_context(ctx, inputs, output, **kwargs): ...

    @staticmethod
    def backward(ctx, *grad_outputs): ...

# API Functions 
class Wrapper(TorchModule):
  def __init__(self, functions, *args, **kwargs):
    super().__init__()
    self.functions = functions
    self.arg_list = args
    self.kwarg_list = kwargs

  def forward(self, input):
    output = input
    for i, (func, args, kwargs) in enumerate(zip(self.functions, self.arg_list, self.kwarg_list):
      output = func.apply(output, *args[i], **kwargs[i])
      #The important thing to note is that func.apply is not actually 
      #performing any computation in this step but rather it is 
      #accumulating the operations defined by the user using my external library API. 
      #This is because I do not want to launch the kernel for every function call.
    
    #Outside the for loop, once we are finished accumulating, I launch the kernel and let my library perform all of the staged operations and then return back a PyTorch tensor.






i.e for a simple model here:
class ExampleModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(4096, 4096, bias=False),  
            torch.nn.Linear(4096, 4096, bias=False),   
            torch.nn.ReLU(),    
        )

    def forward(self, x):
        return self.layers(x)

#And we wanted the first two linear layers to use my library's API 
#instead of PyTorch (so the user would need to implement a class 
#that inherits ExternalLibraryFunction as well as providing 
#their implementation of forward,  backward, and 
#setup_context. 

#Let's say the user implemented this class
#and called it ExternalLibraryLinear
first_param = model.layers[0].weight.clone().detach().to("device").requires_grad_(True)
second_param = model.layers[1].weight.clone().detach().to("device").requires_grad_(True)

ExampleModel.layers[0] = Wrapper([ExternalLibraryLinear, ExternalLibraryLinear], [[first_param], [second_param]])
del ExampleModel.layers[1]
</code></pre>
<p>This is the general idea. What I’ve noticed is that the forward pass works just fine. But the backwards pass does not call the custom autograd backwards function the user implements. It seems this is due to the fact that the forward pass returns a staged computation. I think I need to ensure that the <code>Function.apply</code> call creates a proper node in the computational graph, even if the actual computation is deferred. If I instead made it s.t the user forward would directly return the output and we did not do any operation accumulation, then the backwards pass works fine.</p>
<p>The only solution that immediately come to mind is to represent the entire staged computation as a single node in the autograd graph (Just have a shared <code>ExternalLibraryFunction</code> class. Or maybe creating a placeholder tensor that we could directly return from <code>function.apply</code> to ‘trick’ autograd into creating a node.</p>