Repository: pytorch
Issue: torch.func.jacrev fails if model contains full_backward_hook
Labels: module: autograd,triaged,module: functorch

### 🐛 Describe the bug

Hi All,

**TL;DR** :bug: If your model involves using a `full_backward_hook`, computing derivatives throws a missing `setup_context` error. As a full backward hook is technically a `torch.autograd.Function`, and hence requires a `setup_context` method in pytorch2, which it doesn't seem to have by default. 

Here's a minimal reproducible example,
```
import torch
from torch import nn, Tensor
from torch.func import vmap, jacrev, functional_call

from typing import Tuple
from collections import defaultdict

class model(nn.Module):

  def __init__(self, num_input, num_hidden):
    super(model, self).__init__()
    
    self.fc1 = nn.Linear(num_input, num_hidden)
    self.fc2 = nn.Linear(num_hidden, 1)
    
    self.act_func = nn.Tanh()
    
  def forward(self, x):
    x = self.fc1(x)
    x = self.act_func(x)
    x = self.fc2(x)
    return x
    
num_samples = 4096
num_input=2
num_hidden=64

device=torch.device("cpu")

state = defaultdict(dict) #equivalent to optim.state

def forward_pre_hook(module: nn.Module, input: Tuple[Tensor]) -> None:
  a=input[0]
  state[module]['a'] = a

def full_backward_hook(module: nn.Module, grad_input: Tuple[Tensor], grad_output: Tuple[Tensor]) -> None:
  e = grad_output[0]
  state[module]['e'] = e 

#The input
x = torch.randn(num_samples, num_input, device=device)

#Our model
net = model(num_input=num_input,
            num_hidden=num_hidden)

#Add the hooks
for mod in net.modules():
  if(mod.__class__.__name__ == 'Linear'): #avoids registering hook to nn.Tanh() etc.
    mod.register_forward_pre_hook(forward_pre_hook)
    mod.register_full_backward_hook(full_backward_hook)

y = net(x) #compute output (works fine)

#Compute trace of Hessian
def calc_hessian_trace(params, x):

  def output(params, x):
    return functional_call(net, params, x)

  _hessian = jacrev(jacrev(output, argnums=(1)), argnums=(1))(params, x)
  return _hessian.diagonal(0,-2,-1).sum(-1)

laplacian = vmap(calc_hessian_trace, in_dims=(None, 0))(dict(net.named_parameters()), x) #fails 
```

The resultant error (with complete stack trace) is,
```
Traceback (most recent call last):
  File "~/Downloads/hooks_issue.py", line 64, in <module>
    laplacian = vmap(calc_hessian_trace, in_dims=(None, 0))(dict(net.named_parameters()), x) #fails 
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 434, in wrapped
    return _flat_vmap(
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 619, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "~/Downloads/hooks_issue.py", line 61, in calc_hessian_trace
    _hessian = jacrev(jacrev(output, argnums=(1)), argnums=(1))(params, x)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 489, in wrapper_fn
    vjp_out = _vjp_with_argnums(func, *args, argnums=argnums, has_aux=has_aux)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 291, in _vjp_with_argnums
    primals_out = func(*primals)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 489, in wrapper_fn
    vjp_out = _vjp_with_argnums(func, *args, argnums=argnums, has_aux=has_aux)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 291, in _vjp_with_argnums
    primals_out = func(*primals)
  File "~/Downloads/hooks_issue.py", line 59, in output
    return functional_call(net, params, x)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/_functorch/functional_call.py", line 143, in functional_call
    return nn.utils.stateless._functional_call(
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/nn/utils/stateless.py", line 262, in _functional_call
    return module(*args, **kwargs)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "~/Downloads/hooks_issue.py", line 19, in forward
    x = self.fc1(x)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1536, in _call_impl
    args = bw_hook.setup_input_hook(args)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/utils/hooks.py", line 191, in setup_input_hook
    res, input_idx = self._apply_on_tensors(fn, args)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/utils/hooks.py", line 170, in _apply_on_tensors
    new_tensors = torch.nn.modules._functions.BackwardHookFunction.apply(*tensors)
  File "~/anaconda3/envs/pytorch2/lib/python3.10/site-packages/torch/autograd/function.py", line 509, in apply
    raise RuntimeError(
RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. For more details, please see https://pytorch.org/docs/master/notes/extending.func.html
```

I _think_ this emerges because `full_backward_hook` is technically a `torch.autograd.Function` and so when you call any `torch.func` method it requires a `setup_context` in order to handle any outputs within pytorch2.0. In previous versions of pytorch, `full_backward_hook` methods were skipped entirely if I recall correctly. 

Is there a way to add a `setup_context` to a `full_backward_hook`, or perhaps define a custom `full_backward_hook` with a `setup_context` method itself?

After following through the stack trace, I believe I have found where the error is emerging. In https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/_functions.py, `BackwardHookFunction` is defined as a `torch.autograd.Function` object, but it's in the style of pytorch1.x. 

```
class BackwardHookFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, *args):
        ctx.mark_non_differentiable(*[arg for arg in args if not arg.requires_grad])
        return args

    @staticmethod
    def backward(ctx, *args):
        return args
```

So, I'd assume that changing this function to something like,
```
class BackwardHookFunction(torch.autograd.Function):
    @staticmethod
    def forward(*args):
        return args
    
    @staticmethod
    def setup_context(ctx, *args):
        ctx.mark_non_differentiable(*[arg for arg in args if not arg.requires_grad])

    @staticmethod
    def backward(ctx, *args):
        return args
```
Might resolve this issue? However, I have no idea if this would mess with other parts of pytorch. 



### Versions

Collecting environment information...
PyTorch version: 2.0.0
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.10.10 (main, Mar 21 2023, 18:45:11) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-69-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 10.1.243
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2060
Nvidia driver version: 520.61.05
cuDNN version: Probably one of the following:
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn.so.8.3.2
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.3.2
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.3.2
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.3.2
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.3.2
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.3.2
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.3.2
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   39 bits physical, 48 bits virtual
CPU(s):                          16
On-line CPU(s) list:             0-15
Thread(s) per core:              2
Core(s) per socket:              8
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           165
Model name:                      Intel(R) Core(TM) i7-10700K CPU @ 3.80GHz
Stepping:                        5
CPU MHz:                         3800.000
CPU max MHz:                     5100.0000
CPU min MHz:                     800.0000
BogoMIPS:                        7599.80
Virtualisation:                  VT-x
L1d cache:                       256 KiB
L1i cache:                       256 KiB
L2 cache:                        2 MiB
L3 cache:                        16 MiB
NUMA node0 CPU(s):               0-15
Vulnerability Itlb multihit:     KVM: Mitigation: VMX disabled
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:          Mitigation; Enhanced IBRS
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Mitigation; Microcode
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp pku ospke md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] functorch==1.13.0a0+8e2f53b
[pip3] numpy==1.23.5
[pip3] torch==2.0.0
[pip3] torchaudio==2.0.0
[pip3] torchvision==0.15.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0           py310h7f8727e_0  
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0  
[conda] mkl_random                1.2.2           py310h00e6091_0  
[conda] numpy                     1.23.5          py310hd5efca6_0  
[conda] numpy-base                1.23.5          py310h8e6c178_0  
[conda] pytorch                   2.0.0           py3.10_cuda11.8_cudnn8.7.0_0    pytorch
[conda] pytorch-cuda              11.8                 h7e8668a_3    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.0.0               py310_cu118    pytorch
[conda] torchtriton               2.0.0                     py310    pytorch
[conda] torchvision               0.15.0              py310_cu118    pytorch

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @Chillee @samdow @kshitij12345 @janeyx99 @soumith