Marking variable not required in backward calculation if it is not needed

[mfkasim1]

🚀 Feature

I wonder if there is a feature in pytorch's autograd.Function where it can detect which variables' grads are required and not required.

For example:

import torch

class Mul(torch.autograd.Function):
    @staticmethod
    def forward(ctx, inp1, inp2):
        ctx.save_for_backward(inp1, inp2)
        return inp1 * inp2
    @staticmethod
    def backward(ctx, grad_out):
        inp1, inp2 = ctx.saved_tensors
        grad_inp1 = None
        grad_inp2 = None
        if inp1.requires_grad:  # and inp1.grad_is_needed:
            print("Calculating grad_inp1")
            grad_inp1 = grad_out * inp2
        if inp2.requires_grad:  # and inp2.grad_is_needed:
            print("Calculating grad_inp2")
            grad_inp2 = grad_out * inp1
        return grad_inp1, grad_inp2

dtype = torch.float64
inp1 = torch.tensor(1.0, dtype=dtype).requires_grad_()
inp2 = torch.tensor(2.0, dtype=dtype).requires_grad_()

out = Mul.apply(inp1, inp2)
# grad_inp2 is not required
gradinp1 = torch.autograd.grad(out, inp1, create_graph=True)
# prints:
# Calculating grad_inp1
# Calculating grad_inp2

Gradient of inp2 is not required in the grad calculation, but there is no way to know if inp2 grad is required or not.
It would be good if there is a function that can be called to see if inp2 grad is required, e.g. inp2.grad_is_needed.

Motivation

This problem comes up when I am doing multiple gradients (higher order grad) where the gradient calculation spends a lot of time in unnecessary part of the graph. Similar to the reason in #39784.

Pitch

This feature would enable optimizing the gradient calculation for large graph, especially for computation that requires higher order gradients.

Alternatives

For this simple case, I can just mark inp2 as not requiring grad, but for this example below, there is no way other than just accepting it (i.e. the fact that it calculates unnecessary gradients).

import torch

class Pow(torch.autograd.Function):
    @staticmethod
    def forward(ctx, base, pow):
        out = base ** pow
        ctx.save_for_backward(base, pow, out)
        return out
    @staticmethod
    def backward(ctx, grad_out):
        base, pow, out = ctx.saved_tensors
        grad_base = None
        grad_pow = None
        if base.requires_grad:  # and base.grad_is_needed:
            print("Calculating grad_base")
            grad_base = grad_out * pow * out / base
        if pow.requires_grad:  # and pow.grad_is_needed:
            print("Calculating grad_pow")
            grad_pow = grad_out * out * torch.log(base)
        return grad_base, grad_pow

dtype = torch.float64
base = torch.tensor(1.0, dtype=dtype).requires_grad_()
pow = torch.tensor(2.0, dtype=dtype).requires_grad_()

out = Pow.apply(base, pow)
gp = torch.autograd.grad(out, pow, create_graph=True)
# calculating grad_pow is not required below
gbgp = torch.autograd.grad(gp, base, create_graph=True)

Additional context

This is similar to issue #39784 where it has been solved since February, and I am using nightly version 1.9.0.dev20210409 to run the example above.

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7

[albanD]

Yes, you have ctx.needs_input_grad that is a list of booleans that will tell you if the input at that index requires grad to be computed or not.
You can see more details in the custom Function doc: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function.backward
Or in the extending Note in the doc.

That is what you want right?

[mfkasim1]

@albanD yes! thanks for this!

[mfkasim1]

@albanD sorry to reopen the issue. It seems like ctx.needs_input_grad does not behave like what I expected.
For the simple example above:

import torch

class Mul(torch.autograd.Function):
    @staticmethod
    def forward(ctx, inp1, inp2):
        ctx.save_for_backward(inp1, inp2)
        return inp1 * inp2
    @staticmethod
    def backward(ctx, grad_out):
        inp1, inp2 = ctx.saved_tensors
        grad_inp1 = None
        grad_inp2 = None
        if ctx.needs_input_grad[0]:
            print("Calculating grad_inp1")
            grad_inp1 = grad_out * inp2
        if ctx.needs_input_grad[1]:
            print("Calculating grad_inp2")
            grad_inp2 = grad_out * inp1
        return grad_inp1, grad_inp2

dtype = torch.float64
inp1 = torch.tensor(1.0, dtype=dtype).requires_grad_()
inp2 = torch.tensor(2.0, dtype=dtype).requires_grad_()

out = Mul.apply(inp1, inp2)
# grad_inp2 is not required
gradinp1 = torch.autograd.grad(out, inp1, create_graph=True)

The calculation of grad_inp2 is still performed in the example above. It seems like ctx.needs_input_grad[1] is False only if inp2.requires_grad==False

pytorch/torch/csrc/autograd/python_function.cpp

Lines 491 to 493 in 4575028

	PyObject* needs_grad = tensor.requires_grad() ? Py_True : Py_False;
	Py_INCREF(needs_grad);
	PyTuple_SET_ITEM(flags.needs_input_grad.get(), i, needs_grad);

[albanD]

Yes, that is expected I am afraid.
This is only aware of the links that are in the graph. Not the subset required when calling autograd.grad.

I am not sure if we have this information at all... We would need to investigate.

[mfkasim1]

@albanD There is a possible fix for this with Python which I wrote in colab: https://colab.research.google.com/drive/1wEgkI4IsUQBXUEtgFa-ai3yOwJnuvOGD?usp=sharing (you can skip the "problem" part).
In short, we need to check if the variable has a dependency w.r.t. the derived variables (i.e. the inputs in torch.autograd.grad).

However, this requires a global variable to store the derived variables and an inefficient way to check the dependency of the variable w.r.t. the derived variables.
Do you have any idea to solve this?
If the first problem can be fixed, I think I can live with it.

[albanD]

The problem is that by design, we never have this information :/
In particular, the knowledge is a Node needs to be executed is local to that Node in the engine. And even if a Node is not going to get executed, you might want to compute the gradients wrt to its outputs (because you want to capture them with autograd.grad for example). We have to do it this way because Tensors don't exist in the graph, only their grad_fn.

As an example when you do torch.autograd.grad(out, inp1), we actually put a capture on inp1.grad_fn inputs even though inp1.grad_fn won't be executed.

The other issue is that the existing flag needs_input_grad is actually set during creation (based on the inputs requires grad). And there is no effect of the current backward run on the content of this flag or any API that exist for the engine to tell the Node that for this particular run, some inputs don't need gradient.

[soulitzer]

Not sure if this would solve your issue but we did have a PR that solves something similar #52180. So there might be a way you could try to do this, but depending your specific graph it might not always work. Lets say you have some a = expensive_calculation(x, y, z). If you don't want to execute a.grad_fn but a.grad_fn is part of your graph i.e. you have b = a + c; out = loss(b), you need to make sure a.grad_fn has a topological_nr lower than any node you wish to obtain the grad of. What does this mean exactly?

You can think of topological number as the longest possible path between your node and any leaf node. So if you are a leaf node, your topo_nr is always 0. This is a valid topological ordering because we have the invariant that any parent's topo_nr is one greater than the max topo_nr of its children.

So to apply this to our previous example, if you did autograd.grad(out, (b,)), a.grad_fn wouldn't execute because we know for sure according to the properties of topo_nr that a.grad_fn.topo_nr < b.grad_fn.topo_nr. But if you also want grad w.r.t. c, you'd need to make sure c.grad_fn.topo_nr > a.grad_fn.top_nr to ensure that a.grad_fn doesn't execute.

[albanD]

My understanding is that there is the extra assumption here that a.grad_fn == b.grad_fn ;)
See that the original example in a custom Function.

[soulitzer]

Oh right, I guess that is unrelated to the feature proposal. Since the motivation is for making higher-order stuff more efficient I figured this might be useful to know though. :P

One thing we could do though is to enhance input_needs_grad with information we gained from topo_nr, i.e., if any of the inputs have topo_nr lower than min_topo_nr have the boolean at that index to be False

[mfkasim1]

@soulitzer thanks for the info. it seems that I can use the topo_nr to make the checking more efficient (i.e. helping solve the problem number 2 from the google colab). Is topo_nr available from the Python's interface?

[albanD]

The most horrible way I can think doing it is:

• During the Node preparation, get the current graph_task that we're running from:

  pytorch/torch/csrc/autograd/engine.cpp

  Line 83 in 0a72904

  static thread_local std::shared_ptr<GraphTask> current_graph_task = nullptr;

• From the graph task, if graph_task.exec_info_ does not exist, that means everything needs to be executed and so do as usual
• Otherwise, for every "next_functions", check in graph_task.exec_info_ for that Node. And for each Node, see if should_execute() is True.
• This will tell you if, for that particular call to autograd.grad(), you need to compute the gradients for it or not. (that gradient will be discarded by the engine here anyways in this case)
• for custom Function, update input_needs_grad for that one run only (thread safety issue here!)
• for regular Nodes, we might want to update the should_compute_output as well. That will allow us to trim graphs for codegen defined functions as well :)

@soulitzer what do you think?

[soulitzer]

@mfkasim1 No, not currently. Should it be though? I'm thinking it should at least be exposed as grad_fn._topo_nr for debugging purposes since how the topo_nr's of the nodes interact does make an impact on performance. @albanD what do you think?

@albanD (edit: realized I completely misread some code lol) updating needs_input_grad is more straightforward than I initially thought. But still don't see any way around that thread safety issue you point out, i.e., if multiple threads try to backward over the same node at the same time wrt different inputs.

[MatthiasKohl]

Has there been any progress on this issue? I'm currently running into this issue as well, and it seems like as of version 1.11, it is still happening the same way as described here