`torch.compile` produces `RuntimeError` on function wrapped with `torch.func.grad`
Labels: triaged, oncall: pt2, module: functorch

### 🐛 Describe the bug

Trying out the new `torch.compile` API and am getting a `RuntimeError: Cannot access data pointer of Tensor that doesn't have storage` on a function that has been wrapped with `torch.func.grad`. Why does this not work?

Minimal replication:

```python
import torch
from torch.func import grad

def f(x):
    return (0.5 * x**2).sum()


x = torch.randn(10)

f(x) # This works
print(torch.allclose(grad(f)(x), x)) # This works

fjit = torch.compile(grad(f))
fjit(x) # Error, see below
```

<details>
  <summary><strong>ERROR DETAILS</strong></summary>
<pre>

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[2], line 14
     11 f(x) # Ok, this works
     12 print(torch.allclose(grad(f)(x), x)) # This works
---> 14 fjit(x) # Error, see below

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:209, in _TorchDynamoContext.__call__.<locals>._fn(*args, **kwargs)
    207 dynamic_ctx.__enter__()
    208 try:
--> 209     return fn(*args, **kwargs)
    210 finally:
    211     set_eval_frame(prior)

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py:1364, in grad.<locals>.wrapper(*args, **kwargs)
   1362 @wraps(func)
   1363 def wrapper(*args, **kwargs):
-> 1364     results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)
   1365     if has_aux:
   1366         grad, (_, aux) = results

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/_functorch/vmap.py:39, in doesnt_support_saved_tensors_hooks.<locals>.fn(*args, **kwargs)
     36 @functools.wraps(f)
     37 def fn(*args, **kwargs):
     38     with torch.autograd.graph.disable_saved_tensors_hooks(message):
---> 39         return f(*args, **kwargs)

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py:1224, in grad_and_value.<locals>.wrapper(*args, **kwargs)
   1222 # See NOTE [grad and vjp interaction with no_grad]
   1223 with torch.enable_grad():
-> 1224     args = _wrap_all_tensors(args, level)
   1225     kwargs = _wrap_all_tensors(kwargs, level)
   1226     diff_args = _slice_argnums(args, argnums, as_tuple=False)

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py:91, in _wrap_all_tensors(tensor_pytree, level)
     90 def _wrap_all_tensors(tensor_pytree, level):
---> 91     return tree_map(partial(_wrap_tensor_for_grad, level=level), tensor_pytree)

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/utils/_pytree.py:195, in tree_map(fn, pytree)
    194 def tree_map(fn: Any, pytree: PyTree) -> PyTree:
--> 195     flat_args, spec = tree_flatten(pytree)
    196     return tree_unflatten([fn(i) for i in flat_args], spec)

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/utils/_pytree.py:158, in tree_flatten(pytree)
    156 children_specs : List['TreeSpec'] = []
    157 for child in child_pytrees:
--> 158     flat, child_spec = tree_flatten(child)
    159     result += flat
    160     children_specs.append(child_spec)

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/utils/_pytree.py:148, in tree_flatten(pytree)
    144 """Flattens a pytree into a list of values and a TreeSpec that can be used
    145 to reconstruct the pytree.
    146 """
    147 if _is_leaf(pytree):
--> 148     return [pytree], LeafSpec()
    150 node_type = _get_node_type(pytree)
    151 flatten_fn = SUPPORTED_NODES[node_type].flatten_fn

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/utils/_pytree.py:137, in LeafSpec.__init__(self)
    136 def __init__(self) -> None:
--> 137     super().__init__(None, None, [])
    138     self.num_leaves = 1

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:330, in catch_errors_wrapper.<locals>.catch_errors(frame, cache_size)
    327             return hijacked_callback(frame, cache_size, hooks)
    329 with compile_lock:
--> 330     return callback(frame, cache_size, hooks)

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:404, in convert_frame.<locals>._convert_frame(frame, cache_size, hooks)
    402 counters["frames"]["total"] += 1
    403 try:
--> 404     result = inner_convert(frame, cache_size, hooks)
    405     counters["frames"]["ok"] += 1
    406     return result

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:109, in wrap_convert_context.<locals>._fn(*args, **kwargs)
    107 torch.random.set_rng_state(rng_state)
    108 if torch.cuda.is_available():
--> 109     torch.cuda.set_rng_state(cuda_rng_state)
    110 torch.fx.graph_module._forward_from_src = prior_fwd_from_src

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/cuda/random.py:64, in set_rng_state(new_state, device)
     61     default_generator = torch.cuda.default_generators[idx]
     62     default_generator.set_state(new_state_copy)
---> 64 _lazy_call(cb)

File /HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/cuda/__init__.py:183, in _lazy_call(callable, **kwargs)
    181 def _lazy_call(callable, **kwargs):
    182     if is_initialized():
--> 183         callable()
    184     else:
    185         # TODO(torch_deploy): this accesses linecache, which attempts to read the
    186         # file system to get traceback info. Patch linecache or do something
    187         # else here if this ends up being important.
    188         global _lazy_seed_tracker

File HOME/miniconda3/envs/et-on-fms/lib/python3.10/site-packages/torch/cuda/random.py:62, in set_rng_state.<locals>.cb()
     60     idx = current_device()
     61 default_generator = torch.cuda.default_generators[idx]
---> 62 default_generator.set_state(new_state_copy)

RuntimeError: Cannot access data pointer of Tensor that doesn't have storage
</pre>
</details>


### Versions

Collecting environment information...
PyTorch version: 2.0.0.dev20230213+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 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.25.0
Libc version: glibc-2.31

Python version: 3.10.9 | packaged by conda-forge | (main, Feb  2 2023, 20:20:04) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-137-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 515.86.01
cuDNN version: Could not collect
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:                   43 bits physical, 48 bits virtual
CPU(s):                          256
On-line CPU(s) list:             0-255
Thread(s) per core:              2
Core(s) per socket:              64
Socket(s):                       2
NUMA node(s):                    8
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           49
Model name:                      AMD EPYC 7742 64-Core Processor
Stepping:                        0
Frequency boost:                 enabled
CPU MHz:                         3397.663
CPU max MHz:                     2250.0000
CPU min MHz:                     1500.0000
BogoMIPS:                        4491.77
Virtualization:                  AMD-V
L1d cache:                       4 MiB
L1i cache:                       4 MiB
L2 cache:                        64 MiB
L3 cache:                        512 MiB
NUMA node0 CPU(s):               0-15,128-143
NUMA node1 CPU(s):               16-31,144-159
NUMA node2 CPU(s):               32-47,160-175
NUMA node3 CPU(s):               48-63,176-191
NUMA node4 CPU(s):               64-79,192-207
NUMA node5 CPU(s):               80-95,208-223
NUMA node6 CPU(s):               96-111,224-239
NUMA node7 CPU(s):               112-127,240-255
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Vulnerable
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; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
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 mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate sme ssbd mba sev ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif umip rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.24.2
[pip3] pytorch-triton==2.0.0+0d7e753227
[pip3] torch==2.0.0.dev20230213+cu117
[conda] numpy                     1.24.2                   pypi_0    pypi
[conda] pytorch-triton            2.0.0+0d7e753227          pypi_0    pypi
[conda] torch                     2.0.0.dev20230213+cu117          pypi_0    pypi

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @zou3519 @Chillee @samdow @kshitij12345 @janeyx99