Repository: pytorch
Issue: torch.func transforms don't yet support saved tensor hooks.
Labels: triaged,module: vmap

### 🚀 The feature, motivation and pitch

Hello pytorch developers, I encountered the following error when I am executing my code: `RuntimeError: torch.func transforms don't yet support saved tensor hooks. Please open an issue with your use case.` Basically I am trying to do gradient checkpointing of a function which contains vmap as one of its operations. The minimum code to reproduce this error is given as follows:
```python
import torch
from torch.utils.checkpoint import checkpoint
A = torch.randn((3, 8, 8), dtype=torch.float64, requires_grad=True)

def get_loss(A):
    ortho_A, _ = torch.func.vmap(torch.linalg.qr)(A)
    return torch.sum(ortho_A)

loss = checkpoint(get_loss, A, use_reentrant=False)
loss.backward()
print(A.grad)
```

The full error message is:
```
Traceback (most recent call last):
  File "/path/to/tests/vmapcheckpoint.py", line 9, in <module>
    loss = checkpoint(get_loss, A, use_reentrant=False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Applications/anaconda3/lib/python3.11/site-packages/torch/_compile.py", line 24, in inner
    return torch._dynamo.disable(fn, recursive)(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Applications/anaconda3/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/Applications/anaconda3/lib/python3.11/site-packages/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/Applications/anaconda3/lib/python3.11/site-packages/torch/utils/checkpoint.py", line 458, in checkpoint
    ret = function(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/to/tests/vmapcheckpoint.py", line 6, in get_loss
    ortho_A, _ = torch.func.vmap(torch.linalg.qr)(A)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Applications/anaconda3/lib/python3.11/site-packages/torch/_functorch/apis.py", line 188, in wrapped
    return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Applications/anaconda3/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 266, in vmap_impl
    return _flat_vmap(
           ^^^^^^^^^^^
  File "/Applications/anaconda3/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 37, in fn
    with torch.autograd.graph.disable_saved_tensors_hooks(message):
  File "/Applications/anaconda3/lib/python3.11/contextlib.py", line 137, in __enter__
    return next(self.gen)
           ^^^^^^^^^^^^^^
  File "/Applications/anaconda3/lib/python3.11/site-packages/torch/autograd/graph.py", line 322, in disable_saved_tensors_hooks
    torch._C._autograd._saved_tensors_hooks_disable(error_message)
RuntimeError: torch.func transforms don't yet support saved tensor hooks. Please open an issue with your use case.
```

I wonder if it is possible to make vmap compatible for gradient checkpointing, which is crucial for my code. Thanks!

### Alternatives

_No response_

### Additional context

_No response_

cc @zou3519