Double-backward with full_backward_hook causes RuntimeError in PyTorch 1.13

[f-dangel]

Hi,

I am using double-backward calls to compute Hessian matrices, in combination with PyTorch's full_backward_hooks. After upgrading from 1.12.1 to 1.13.0, I now run into the following error in the second backward pass:

RuntimeError: Module backward hook for grad_input is called before the grad_output one. This happens because the gradient in your nn.Module flows to the Module's input without passing through the Module's output. Make sure that the output depends on the input and that the loss is computed based on the output.

The following snippet reproduces my problem when I try to compute the Hessian of f(x, y) w.r.t. x where all symbols are scalars for simplicity.

"""Compute the scalar-valued second-order derivative of f(x, y) w.r.t. x.

Use Hessian-vector products (double-backward pass) in combination with
full_backward_hook.
"""

from torch import ones_like, rand, rand_like
from torch.autograd import grad
from torch.nn import MSELoss

x = rand(1)
x.requires_grad_(True)
y = rand_like(x)

# without hook (working in 1.12.1 and 1.13.0)
lossfunc = MSELoss()
f = lossfunc(x, y)

(gradx_f,) = grad(f, x, create_graph=True)
(gradxgradx_f,) = grad(gradx_f @ ones_like(x), x)

# with hook (working in 1.12.1 and broken in 1.13.0
lossfunc = MSELoss()


def hook(module, grad_input, grad_output):
    print("This is a test hook")


lossfunc.register_full_backward_hook(hook)

f = lossfunc(x, y)

# this line triggers the backward hook as expected
(gradx_f,) = grad(f, x, create_graph=True)
# the double-backward with hook crashes in 1.13, but used to work before
try:
    (gradxgradx_f,) = grad(gradx_f @ ones_like(x), x)
except RuntimeError as e:
    print(f"Caught RuntimeError: {e}")

Is this the intended behavior? If so, how do I compute higher-order derivatives through multiple backward calls, while using hooks, e.g. for monitoring?

Best,
Felix

Versions

PyTorch version: 1.13.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04) 11.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.10

Python version: 3.7.6 (default, Jan 8 2020, 19:59:22) [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-52-generic-x86_64-with-debian-bookworm-sid
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] backpack-for-pytorch==1.5.1.dev14+g5401cde6
[pip3] mypy==0.940
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.1
[pip3] pytorch-memlab==0.2.3
[pip3] torch==1.13.0
[pip3] torchvision==0.10.0
[conda] backpack-for-pytorch 1.5.1.dev14+g5401cde6 dev_0
[conda] numpy 1.20.1 pypi_0 pypi
[conda] pytorch-memlab 0.2.3 pypi_0 pypi
[conda] torch 1.13.0 pypi_0 pypi
[conda] torchvision 0.10.0 pypi_0 pypi

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

[soulitzer]

I think the assertion is wrong, and we can just remove it. Previously, we assumed that everytime the backward post hook fires, the backward pre hook must've fired first - and so we rely on the pre hook to setup the grad_outputs information. Double backward fails because gradx_f does not actually depend on the output of f, so grad_outputs information was not setup when the post hook fires..

This was only "working" previously because we failed to clear the grad_output after the hook fires #82788, so those buffers would've been leaked, and because you aren't truly backwarding through that module again, the grad_outputs that you are getting from that hook don't really mean anything.

[albanD]

Interesting! Sounds good!

[gchanan]

should this be in the 1.13.1 milestone?