[functorch] [FakeTensorMode, meta tensor] + aot_autograd Bug.
Labels: triaged, oncall: pt2, module: fakeTensor, module: aotdispatch

### 🐛 Describe the bug

I am trying to use FakeTensor and aot_autograd to capture the computation graph, but I met below errors. Can anyone help me out?

# FakeTensorMode case
In this case, I got errors like `TypeError: Multiple dispatch failed for 'torch._ops.aten.t.default'; all __torch_dispatch__ handlers returned NotImplemented`.


```python
import torch
from torch.nn import Linear
from torchdistx.fake import fake_mode
from torch._subclasses.fake_tensor import FakeTensorMode, FakeTensor
from torch._functorch.aot_autograd import aot_export_joint_simple


class TestModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = Linear(1024, 4096)
        self.linear2 = Linear(4096, 1024)

    def forward(self, x):
        y = self.linear(x)
        z = self.linear2(y)
        # loss = torch.sum(z)
        return tuple([z])


with FakeTensorMode():
    sample_input = torch.randn(4, 512, 1024)
    loss = torch.rand(4, 512, 4096)

    model = TestModel()
    z = model(sample_input)

    graph_module = aot_export_joint_simple(model, tuple([sample_input]), trace_joint=True)
    print(graph_module)

```

```
[2023-06-13 20:10:10,474] torch.fx.experimental.proxy_tensor.__not_implemented: [DEBUG] ProxyTensorMode tensors without proxy had unrecognized subclasses: [<class 'torch._subclasses.fake_tensor.FakeTensor'>]
[2023-06-13 20:10:10,474] torch._subclasses.fake_tensor.__not_implemented: [DEBUG] FakeTensor mode already active: <torch._subclasses.fake_tensor.FakeTensorMode object at 0x7f91b9efbfa0> in [<torch._subclasses.fake_tensor.FakeTensorMode object at 0x7f91b9efbfa0>, <torch.fx.experimental.proxy_tensor.ProxyTorchDispatchMode object at 0x7f91c24e23a0>]
Traceback (most recent call last):
  File "/Users/connolly/Documents/GitHub/Autoplanner/test/fake_tensor_bug_issue.py", line 28, in <module>
    graph_module = aot_export_joint_simple(model, tuple([sample_input]), trace_joint=True)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3960, in aot_export_joint_simple
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 4050, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 180, in time_wrapper
    r = func(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3262, in create_aot_dispatcher_function
    compiled_fn = compiler_fn(flat_fn, fake_flat_args, aot_config, fw_metadata=fw_metadata)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 2083, in aot_wrapper_dedupe
    return compiler_fn(flat_fn, leaf_flat_args, aot_config, fw_metadata=fw_metadata)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 2263, in aot_wrapper_synthetic_base
    return compiler_fn(flat_fn, flat_args, aot_config, fw_metadata=fw_metadata)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1500, in aot_dispatch_base_graph
    fw_module = create_functionalized_graph(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1387, in create_functionalized_graph
    fx_g = make_fx(helper, decomposition_table=aot_config.decompositions)(*args)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 769, in wrapped
    t = dispatch_trace(wrap_key(func, args, fx_tracer, pre_dispatch), tracer=fx_tracer, concrete_args=tuple(phs))
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 463, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py", line 810, in trace
    (self.create_arg(fn(*args)),),
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 480, in wrapped
    out = f(*tensors)
  File "<string>", line 1, in <lambda>
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1379, in fwd_helper
    return functionalized_f_helper(*args)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1329, in functionalized_f_helper
    f_outs = fn(*f_args)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1128, in inner_fn
    outs = fn(*args)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3374, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py", line 788, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 429, in call_module
    return forward(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py", line 781, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/connolly/Documents/GitHub/Autoplanner/test/fake_tensor_bug_issue.py", line 15, in forward
    y = self.linear(x)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py", line 788, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 429, in call_module
    return forward(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py", line 781, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
TypeError: Multiple dispatch failed for 'torch._ops.aten.t.default'; all __torch_dispatch__ handlers returned NotImplemented:

  - mode object <torch.fx.experimental.proxy_tensor.ProxyTorchDispatchMode object at 0x7f91c24e23a0>
  - tensor subclass <class 'torch._subclasses.fake_tensor.FakeTensor'>
```

# torchdistx fake_mode case
If I replace the `with FakeTensorMode():` by `with fake_mode()` in torchdistx, it gets below errors:
```
Traceback (most recent call last):
  File "/Users/connolly/Documents/GitHub/Autoplanner/test/fake_tensor_bug_issue.py", line 28, in <module>
    graph_module = aot_export_joint_simple(model, tuple([sample_input]), trace_joint=True)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3960, in aot_export_joint_simple
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 4050, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 180, in time_wrapper
    r = func(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3201, in create_aot_dispatcher_function
    fake_flat_args = process_inputs(flat_args)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3199, in process_inputs
    return [convert(idx, x) for idx, x in enumerate(flat_args)]
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3199, in <listcomp>
    return [convert(idx, x) for idx, x in enumerate(flat_args)]
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3197, in convert
    return fake_mode.from_tensor(x, static_shapes=False)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 1588, in from_tensor
    return self.fake_tensor_converter(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 341, in __call__
    return self.from_real_tensor(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 294, in from_real_tensor
    out = self.meta_converter(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_subclasses/meta_utils.py", line 531, in __call__
    r = self.meta_tensor(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_subclasses/meta_utils.py", line 410, in meta_tensor
    s = t.untyped_storage()
NotImplementedError: Cannot access storage of FakeTensorImpl
```

# Meta Tensor case
In meta tensor case, below code works for small models. But for some complicated models like Bert, it raises Errors shown below.

```python
from transformers import BertModel, BertConfig
with torch.device("meta"):
    sample_input = torch.randint(0, 30522, [4, 512])

    model = BertModel(BertConfig())
    z = model(sample_input)

    graph_module = aot_export_module(model, tuple([sample_input]),output_loss_index=0, trace_joint=True)
    print(graph_module)
```

```
Traceback (most recent call last):
  File "/Users/connolly/Documents/GitHub/Autoplanner/test/fake_tensor_bug_issue.py", line 29, in <module>
    graph_module = aot_export_module(model, tuple([sample_input]),output_loss_index=0, trace_joint=True)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3868, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 4050, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 180, in time_wrapper
    r = func(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3212, in create_aot_dispatcher_function
    fw_metadata = run_functionalized_fw_and_collect_metadata(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 733, in inner
    flat_f_outs = f(*flat_f_args)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3374, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3822, in fn_to_trace
    out = functional_call(*args)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 3347, in functional_call
    out = mod(*args[params_len:], **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/transformers/models/bert/modeling_bert.py", line 1013, in forward
    embedding_output = self.embeddings(
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/transformers/models/bert/modeling_bert.py", line 238, in forward
    embeddings = self.dropout(embeddings)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/modules/dropout.py", line 59, in forward
    return F.dropout(input, self.p, self.training, self.inplace)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/functional.py", line 1267, in dropout
    return handle_torch_function(dropout, (input,), input, p=p, training=training, inplace=inplace)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/overrides.py", line 1541, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/utils/_device.py", line 76, in __torch_function__
    return func(*args, **kwargs)
  File "/Users/connolly/opt/anaconda3/envs/astropy/lib/python3.9/site-packages/torch/nn/functional.py", line 1270, in dropout
    return _VF.dropout_(input, p, training) if inplace else _VF.dropout(input, p, training)
RuntimeError: 0 INTERNAL ASSERT FAILED at "/Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/core/boxing/KernelFunction.cpp":19, please report a bug to PyTorch. fallthrough_kernel was executed but it should have been short-circuited by the dispatcher. This could occur if you registered a fallthrough kernel as a override for a specific operator (as opposed to a backend fallback); this is NOT currently supported, and we do not intend to add support for it in the near future.  If you do find yourself in need of this, let us know in the bug tracker.
```


### Versions

PyTorch version: 2.1.0.dev20230612
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.4 (x86_64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: version 3.24.0
Libc version: N/A

Python version: 3.9.13 (main, Aug 25 2022, 18:29:29)  [Clang 12.0.0 ] (64-bit runtime)
Python platform: macOS-10.16-x86_64-i386-64bit
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

CPU:
Intel(R) Core(TM) i7-8850H CPU @ 2.60GHz

Versions of relevant libraries:
[pip3] functorch==0.2.1
[pip3] numpy==1.23.2
[pip3] torch==2.1.0.dev20230612
[pip3] torchdistx==0.3.0.dev0+cpu
[conda] functorch                 0.2.1                    pypi_0    pypi
[conda] numpy                     1.23.2                   pypi_0    pypi
[conda] torch                     2.1.0.dev20230612          pypi_0    pypi
[conda] torchdistx                0.3.0.dev0+cpu           pypi_0    pypi


cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305