torch.load fails under FakeTensorMode for GPT2 model
Labels: triaged, ezyang's list, oncall: pt2, module: fakeTensor, module: dynamo, release notes: dynamo

### 🐛 Describe the bug

I will break up this issue in 2 parts. The first discusses the issue without any modification to PyTorch code. The second one hacks `torch._util._rebuild_tensor` and tries to export the model to FX without success. I did this because I am not sure if the second part is a side effect of a bad workaround or an entirely different issue. If it is the latter, I am ok in filing a second issue for it after confirmation.

#### Part 1: Before the hack

Model loading works when called outside `FakeTensorMode` context, but it fails when called within it.
This becomes relevant after https://github.com/pytorch/pytorch/pull/100017/ in which we can fakefy input and model parameters before calling `torch._dynamo.export(fake_model, *fake_args, fake_mode=fake_mode, **fake_kwargs)` with the pre-instantiated `FakeTensorMode`.

**Repro:**

```python
from torch._subclasses import fake_tensor
from torch.fx.experimental.symbolic_shapes import ShapeEnv
import transformers

fake_mode = fake_tensor.FakeTensorMode(
    allow_non_fake_inputs=False,
    shape_env=ShapeEnv(
        allow_scalar_outputs=False, allow_dynamic_output_shape_ops=False
    ),
)

fake_model = transformers.AutoModel.from_pretrained("sshleifer/tiny-gpt2")
assert fake_model is not None
with fake_mode:
    fake_model = transformers.AutoModel.from_pretrained("sshleifer/tiny-gpt2")  # raises OSError: Unable to load weights from pytorch checkpoint file for '...' at  If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True.
```

This is the full error:
```
Some weights of the model checkpoint at sshleifer/tiny-gpt2 were not used when initializing GPT2Model: ['lm_head.weight']
- This IS expected if you are initializing GPT2Model from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing GPT2Model from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/modeling_utils.py:463 in           │
│ load_state_dict                                                                                  │
│                                                                                                  │
│    460 │   │   │   )                                                                             │
│    461 │   │   return safe_load_file(checkpoint_file)                                            │
│    462 │   try:                                                                                  │
│ ❱  463 │   │   return torch.load(checkpoint_file, map_location="cpu")                            │
│    464 │   except Exception as e:                                                                │
│    465 │   │   try:                                                                              │
│    466 │   │   │   with open(checkpoint_file) as f:                                              │
│                                                                                                  │
│ /opt/pytorch/torch/serialization.py:1030 in load                                                 │
│                                                                                                  │
│   1027 │   │   │   │   return _legacy_load(opened_file, map_location, _weights_only_unpickler,   │
│   1028 │   │   │   except RuntimeError as e:                                                     │
│   1029 │   │   │   │   raise pickle.UnpicklingError(UNSAFE_MESSAGE + str(e)) from None           │
│ ❱ 1030 │   │   return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args  │
│   1031                                                                                           │
│   1032                                                                                           │
│   1033 # Register pickling support for layout instances such as                                  │
│                                                                                                  │
│ /opt/pytorch/torch/serialization.py:1258 in _legacy_load                                         │
│                                                                                                  │
│   1255 │   _sys_info = pickle_module.load(f, **pickle_load_args)                                 │
│   1256 │   unpickler = UnpicklerWrapper(f, **pickle_load_args)                                   │
│   1257 │   unpickler.persistent_load = persistent_load                                           │
│ ❱ 1258 │   result = unpickler.load()                                                             │
│   1259 │                                                                                         │
│   1260 │   deserialized_storage_keys = pickle_module.load(f, **pickle_load_args)                 │
│   1261                                                                                           │
│                                                                                                  │
│ /opt/pytorch/torch/_utils.py:201 in _rebuild_tensor_v2                                           │
│                                                                                                  │
│   198 def _rebuild_tensor_v2(                                                                    │
│   199 │   storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata=None    │
│   200 ):                                                                                         │
│ ❱ 201 │   tensor = _rebuild_tensor(storage, storage_offset, size, stride)                        │
│   202 │   tensor.requires_grad = requires_grad                                                   │
│   203 │   if metadata:                                                                           │
│   204 │   │   set_tensor_metadata(tensor, metadata)                                              │
│                                                                                                  │
│ /opt/pytorch/torch/_utils.py:180 in _rebuild_tensor                                              │
│                                                                                                  │
│   177 def _rebuild_tensor(storage, storage_offset, size, stride):                                │
│   178 │   # first construct a tensor with the correct dtype/device                               │
│   179 │   t = torch.tensor([], dtype=storage.dtype, device=storage._untyped_storage.device)      │
│ ❱ 180 │   return t.set_(storage._untyped_storage, storage_offset, size, stride)                  │
│   181                                                                                            │
│   182                                                                                            │
│   183 def get_tensor_metadata(tensor):                                                           │
│                                                                                                  │
│ /opt/pytorch/torch/utils/_stats.py:20 in wrapper                                                 │
│                                                                                                  │
│   17 │   │   if fn.__qualname__ not in simple_call_counter:                                      │
│   18 │   │   │   simple_call_counter[fn.__qualname__] = 0                                        │
│   19 │   │   simple_call_counter[fn.__qualname__] = simple_call_counter[fn.__qualname__] + 1     │
│ ❱ 20 │   │   return fn(*args, **kwargs)                                                          │
│   21 │   return wrapper                                                                          │
│   22                                                                                             │
│                                                                                                  │
│ /opt/pytorch/torch/_subclasses/fake_tensor.py:1160 in __torch_dispatch__                         │
│                                                                                                  │
│   1157 │   def __torch_dispatch__(self, func, types, args=(), kwargs=None):                      │
│   1158 │   │   assert self not in _get_current_dispatch_mode_stack(), func                       │
│   1159 │   │   try:                                                                              │
│ ❱ 1160 │   │   │   return self.dispatch(func, types, args, kwargs)                               │
│   1161 │   │   except TypeError:                                                                 │
│   1162 │   │   │   log.exception("fake tensor raised TypeError")                                 │
│   1163 │   │   │   raise                                                                         │
│                                                                                                  │
│ /opt/pytorch/torch/_subclasses/fake_tensor.py:1318 in dispatch                                   │
│                                                                                                  │
│   1315 │   │                                                                                     │
│   1316 │   │   # we are falling through to running non constant tensors, any input constant tha  │
│   1317 │   │   # is written to must be invalidated                                               │
│ ❱ 1318 │   │   self.invalidate_written_to_constants(func, flat_arg_fake_tensors, args, kwargs)   │
│   1319 │   │                                                                                     │
│   1320 │   │   # Try for fastpath                                                                │
│   1321 │   │   if has_symbolic_sizes:                                                            │
│                                                                                                  │
│ /opt/pytorch/torch/_subclasses/fake_tensor.py:1557 in invalidate_written_to_constants            │
│                                                                                                  │
│   1554 │   │   any_constant = any(e.constant is not None for e in flat_arg_fake_tensors)         │
│   1555 │   │   if any_constant and get_schema_info(func).is_mutable():                           │
│   1556 │   │   │   schema_info = get_schema_info(func)                                           │
│ ❱ 1557 │   │   │   _, new_kwargs = normalize_function(                                           │
│   1558 │   │   │   │   func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True         │
│   1559 │   │   │   )                                                                             │
│   1560 │   │   │   for k, v in new_kwargs.items():                                               │
│                                                                                                  │
│ /opt/pytorch/torch/fx/operator_schemas.py:297 in normalize_function                              │
│                                                                                                  │
│   294 │   │   new_args_and_kwargs = _args_kwargs_to_normalized_args_kwargs(sig, args, kwargs,    │
│   295 │   else:                                                                                  │
│   296 │   │   assert callable(target)                                                            │
│ ❱ 297 │   │   torch_op_schemas = get_signature_for_torch_op(target)                              │
│   298 │   │   matched_schemas = []                                                               │
│   299 │   │   if torch_op_schemas:                                                               │
│   300 │   │   │   # Iterate through all of the schema until we find one that matches             │
│                                                                                                  │
│ /opt/pytorch/torch/fx/operator_schemas.py:167 in get_signature_for_torch_op                      │
│                                                                                                  │
│   164 │   │   │   return (None, None) if return_schemas else None                                │
│   165 │   │   schemas = torch._C._jit_get_schemas_for_operator(aten_fn)                          │
│   166 │                                                                                          │
│ ❱ 167 │   signatures = [_torchscript_schema_to_signature(schema) for schema in schemas]          │
│   168 │   return (signatures, schemas) if return_schemas else signatures                         │
│   169                                                                                            │
│   170 @compatibility(is_backward_compatible=False)                                               │
│                                                                                                  │
│ /opt/pytorch/torch/fx/operator_schemas.py:167 in <listcomp>                                      │
│                                                                                                  │
│   164 │   │   │   return (None, None) if return_schemas else None                                │
│   165 │   │   schemas = torch._C._jit_get_schemas_for_operator(aten_fn)                          │
│   166 │                                                                                          │
│ ❱ 167 │   signatures = [_torchscript_schema_to_signature(schema) for schema in schemas]          │
│   168 │   return (signatures, schemas) if return_schemas else signatures                         │
│   169                                                                                            │
│   170 @compatibility(is_backward_compatible=False)                                               │
│                                                                                                  │
│ /opt/pytorch/torch/fx/operator_schemas.py:70 in _torchscript_schema_to_signature                 │
│                                                                                                  │
│    67 │   from inspect import Parameter                                                          │
│    68 │   parameters : List[Parameter] = []                                                      │
│    69 │   for arg in ts_schema.arguments:                                                        │
│ ❱  70 │   │   arg_type = _torchscript_type_to_python_type(arg.type)                              │
│    71 │   │   default = arg.default_value if arg.has_default_value() else Parameter.empty        │
│    72 │   │   # TODO: Figure out if this is safe. It seems like when generating the type signa   │
│    73 │   │   # PythonArgParser, we emit signatures with `input` instead of `self` as the firs   │
│                                                                                                  │
│ /opt/pytorch/torch/fx/operator_schemas.py:64 in _torchscript_type_to_python_type                 │
│                                                                                                  │
│    61 │   eval'ing the annotation_str. _type_eval_globals sets up expressions                    │
│    62 │   like "List" and "Future" to map to actual types (typing.List and jit.Future)           │
│    63 │   """                                                                                    │
│ ❱  64 │   return eval(ts_type.annotation_str, _type_eval_globals)                                │
│    65                                                                                            │
│    66 def _torchscript_schema_to_signature(ts_schema : torch._C.FunctionSchema) -> inspect.Sig   │
│    67 │   from inspect import Parameter                                                          │
│ <string>:1 in <module>                                                                           │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
NameError: name 'Storage' is not defined

During handling of the above exception, another exception occurred:

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/modeling_utils.py:467 in           │
│ load_state_dict                                                                                  │
│                                                                                                  │
│    464 │   except Exception as e:                                                                │
│    465 │   │   try:                                                                              │
│    466 │   │   │   with open(checkpoint_file) as f:                                              │
│ ❱  467 │   │   │   │   if f.read(7) == "version":                                                │
│    468 │   │   │   │   │   raise OSError(                                                        │
│    469 │   │   │   │   │   │   "You seem to have cloned a repository without having git-lfs ins  │
│    470 │   │   │   │   │   │   "git-lfs and run `git lfs install` followed by `git lfs pull` in  │
│                                                                                                  │
│ /opt/conda/envs/ptca/lib/python3.8/codecs.py:322 in decode                                       │
│                                                                                                  │
│    319 │   def decode(self, input, final=False):                                                 │
│    320 │   │   # decode input (taking the buffer into account)                                   │
│    321 │   │   data = self.buffer + input                                                        │
│ ❱  322 │   │   (result, consumed) = self._buffer_decode(data, self.errors, final)                │
│    323 │   │   # keep undecoded input until the next call                                        │
│    324 │   │   self.buffer = data[consumed:]                                                     │
│    325 │   │   return result                                                                     │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
UnicodeDecodeError: 'utf-8' codec can't decode byte 0x80 in position 0: invalid start byte

During handling of the above exception, another exception occurred:

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /opt/pytorch/bug_repro.py:16 in <module>                                                         │
│                                                                                                  │
│   13 fake_model = transformers.AutoModel.from_pretrained("sshleifer/tiny-gpt2")                  │
│   14 assert fake_model is not None                                                               │
│   15 with fake_mode:                                                                             │
│ ❱ 16 │   fake_model = transformers.AutoModel.from_pretrained("sshleifer/tiny-gpt2")  # raises    │
│                                                                                                  │
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/models/auto/auto_factory.py:484 in │
│ from_pretrained                                                                                  │
│                                                                                                  │
│   481 │   │   │   )                                                                              │
│   482 │   │   elif type(config) in cls._model_mapping.keys():                                    │
│   483 │   │   │   model_class = _get_model_class(config, cls._model_mapping)                     │
│ ❱ 484 │   │   │   return model_class.from_pretrained(                                            │
│   485 │   │   │   │   pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs,   │
│   486 │   │   │   )                                                                              │
│   487 │   │   raise ValueError(                                                                  │
│                                                                                                  │
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/modeling_utils.py:2604 in          │
│ from_pretrained                                                                                  │
│                                                                                                  │
│   2601 │   │   if from_pt:                                                                       │
│   2602 │   │   │   if not is_sharded and state_dict is None:                                     │
│   2603 │   │   │   │   # Time to load the checkpoint                                             │
│ ❱ 2604 │   │   │   │   state_dict = load_state_dict(resolved_archive_file)                       │
│   2605 │   │   │                                                                                 │
│   2606 │   │   │   # set dtype to instantiate the model under:                                   │
│   2607 │   │   │   # 1. If torch_dtype is not None, we use that dtype                            │
│                                                                                                  │
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/modeling_utils.py:479 in           │
│ load_state_dict                                                                                  │
│                                                                                                  │
│    476 │   │   │   │   │   │   "model. Make sure you have saved the model properly."             │
│    477 │   │   │   │   │   ) from e                                                              │
│    478 │   │   except (UnicodeDecodeError, ValueError):                                          │
│ ❱  479 │   │   │   raise OSError(                                                                │
│    480 │   │   │   │   f"Unable to load weights from pytorch checkpoint file for '{checkpoint_f  │
│    481 │   │   │   │   f"at '{checkpoint_file}'. "                                               │
│    482 │   │   │   │   "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please s  │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
OSError: Unable to load weights from pytorch checkpoint file for '/root/.cache/huggingface/hub/models--sshleifer--tiny-gpt2/snapshots/5f91d94bd9cd7190a9f3216ff93cd1dd95f2c7be/pytorch_model.bin' at 
'/root/.cache/huggingface/hub/models--sshleifer--tiny-gpt2/snapshots/5f91d94bd9cd7190a9f3216ff93cd1dd95f2c7be/pytorch_model.bin'. If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set 
from_tf=True.
```

By patching `torch._utils._rebuild_tensor` with a context manager like

** Repro **

```python

class PytorchPatcher:
    def __init__(self):

        def torch__util__rebuild_tensor_wrapper(storage, storage_offset, size, stride):
            from torch._subclasses.fake_tensor import FakeTensorMode
            from torch.utils._mode_utils import no_dispatch
            from torch.utils._python_dispatch import _get_current_dispatch_mode

            def _rebuild_real_tensor(storage, storage_offset, size, stride):
                t = torch.tensor(
                    [], dtype=storage.dtype, device=storage._untyped_storage.device
                )
                return t.set_(storage._untyped_storage, storage_offset, size, stride)

            mode = _get_current_dispatch_mode()
            if isinstance(mode, FakeTensorMode):
                # Create a real tensor and then convert it to FakeTensor.
                # We cannot directly create a FakeTensor because it tensor.set_(...)
                # is not supported in FakeTensorMode dispatcher.
                with no_dispatch():
                    t = _rebuild_real_tensor(storage, storage_offset, size, stride)
                return mode.from_tensor(t)
            return _rebuild_real_tensor(storage, storage_offset, size, stride)

        # Original version of torch.load.
        self.torch__util_rebuild_tensor = torch._utils._rebuild_tensor

        # Wrapper or modified version of torch functions.
        self.torch__util_rebuild_tensor_wrapper = torch__util__rebuild_tensor_wrapper

    def __enter__(self):
        torch._utils._rebuild_tensor = self.torch__util_rebuild_tensor_wrapper

    def __exit__(self, exc_type, exc_value, traceback):
        torch._utils._rebuild_tensor = self.torch__util_rebuild_tensor
```  

it does work, but I wonder if this is the right way of doing this. My expectation was that under a fake mode context, real tensors would never be created so having to hack `torch._util._rebuild_tensor` was unexpected.

```
torch_patcher = PytorchPatcher()
with fake_mode, torch_patcher:
    fake_model = transformers.AutoModel.from_pretrained("sshleifer/tiny-gpt2")  # raises OSError: Unable to load weights from 
``` 
 
#### Part 2: Exporting to FX after the hack

With the hack above, going a step further and trying to call `torch._dynamo.export` fails because it seems non-fake tensors are found within the fake_model instance. 

I am not sure whether this is a side-effect of patching `torch._util._rebuild_tensor` or an entirely different issue uncovered after a succesful hack.

** Repro **

```python
def create_args():
    tokenizer = transformers.AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2")
    kwargs = tokenizer("Hello world!", return_tensors="pt")
    input_ids = kwargs["input_ids"]
    attention_mask = kwargs["attention_mask"]
    return input_ids, None, attention_mask

def create_kwargs():
    return {"return_dict": False}

converter = fake_mode.fake_tensor_converter
fake_mode.validate_and_convert_non_fake_tensors(fake_model, converter, fake_args, fake_kwargs)  # passes, only fake stuff here

torch._dynamo.export(fake_model, *fake_args, tracing_mode="fake", fake_mode=fake_mode, **fake_kwargs)  # validate_and_convert_non_fake_tensors fails here. what is the difference?
```

That is the error from the export attempt:

```bash
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /opt/pytorch/bug_repro.py:120 in <module>                                                        │
│                                                                                                  │
│   117                                                                                            │
│   118                                                                                            │
│   119                                                                                            │
│ ❱ 120 torch._dynamo.export(fake_model, *create_args(), tracing_mode="fake", fake_mode=fake_mod   │
│                                                                                                  │
│ /opt/pytorch/torch/_dynamo/eval_frame.py:948 in export                                           │
│                                                                                                  │
│    945 │   │   )(f)                                                                              │
│    946 │   │   # TODO(voz): We may have instances of `f` that mutate inputs, we should track si  │
│    947 │   │   try:                                                                              │
│ ❱  948 │   │   │   result_traced = opt_f(*args, **kwargs)                                        │
│    949 │   │   except ConstraintViolationError as e:                                             │
│    950 │   │   │   constraint_violation_error = e                                                │
│    951 │   remove_from_cache(f)                                                                  │
│                                                                                                  │
│ /opt/pytorch/torch/nn/modules/module.py:1522 in _wrapped_call_impl                               │
│                                                                                                  │
│   1519 │   │   if self._compiled_call_impl is not None:                                          │
│   1520 │   │   │   return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]        │
│   1521 │   │   else:                                                                             │
│ ❱ 1522 │   │   │   return self._call_impl(*args, **kwargs)                                       │
│   1523 │                                                                                         │
│   1524 │   def _call_impl(self, *args, **kwargs):                                                │
│   1525 │   │   forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.fo  │
│                                                                                                  │
│ /opt/pytorch/torch/nn/modules/module.py:1531 in _call_impl                                       │
│                                                                                                  │
│   1528 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1529 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1530 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1531 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1532 │   │                                                                                     │
│   1533 │   │   try:                                                                              │
│   1534 │   │   │   result = None                                                                 │
│                                                                                                  │
│ /opt/pytorch/torch/_dynamo/eval_frame.py:294 in _fn                                              │
│                                                                                                  │
│    291 │   │   │   dynamic_ctx = enable_dynamic(self.dynamic, self.export)                       │
│    292 │   │   │   dynamic_ctx.__enter__()                                                       │
│    293 │   │   │   try:                                                                          │
│ ❱  294 │   │   │   │   return fn(*args, **kwargs)                                                │
│    295 │   │   │   finally:                                                                      │
│    296 │   │   │   │   set_eval_frame(prior)                                                     │
│    297 │   │   │   │   dynamic_ctx.__exit__(None, None, None)                                    │
│                                                                                                  │
│ /opt/pytorch/torch/nn/modules/module.py:1522 in _wrapped_call_impl                               │
│                                                                                                  │
│   1519 │   │   if self._compiled_call_impl is not None:                                          │
│   1520 │   │   │   return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]        │
│   1521 │   │   else:                                                                             │
│ ❱ 1522 │   │   │   return self._call_impl(*args, **kwargs)                                       │
│   1523 │                                                                                         │
│   1524 │   def _call_impl(self, *args, **kwargs):                                                │
│   1525 │   │   forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.fo  │
│                                                                                                  │
│ /opt/pytorch/torch/nn/modules/module.py:1531 in _call_impl                                       │
│                                                                                                  │
│   1528 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1529 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1530 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1531 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1532 │   │                                                                                     │
│   1533 │   │   try:                                                                              │
│   1534 │   │   │   result = None                                                                 │
│                                                                                                  │
│ /opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py:751   │
│ in forward                                                                                       │
│                                                                                                  │
│    748 │   │   for layer, heads in heads_to_prune.items():                                       │
│    749 │   │   │   self.h[layer].attn.prune_heads(heads)                                         │
│    750 │                                                                                         │
│ ❱  751 │   @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)                         │
│    752 │   @add_code_sample_docstrings(                                                          │
│    753 │   │   checkpoint=_CHECKPOINT_FOR_DOC,                                                   │
│    754 │   │   output_type=BaseModelOutputWithPastAndCrossAttentions,                            │
│                                                                                                  │
│ /opt/pytorch/torch/_dynamo/eval_frame.py:294 in _fn                                              │
│                                                                                                  │
│    291 │   │   │   dynamic_ctx = enable_dynamic(self.dynamic, self.export)                       │
│    292 │   │   │   dynamic_ctx.__enter__()                                                       │
│    293 │   │   │   try:                                                                          │
│ ❱  294 │   │   │   │   return fn(*args, **kwargs)                                                │
│    295 │   │   │   finally:                                                                      │
│    296 │   │   │   │   set_eval_frame(prior)                                                     │
│    297 │   │   │   │   dynamic_ctx.__exit__(None, None, None)                                    │
│                                                                                                  │
│ /opt/pytorch/torch/_dynamo/external_utils.py:17 in inner                                         │
│                                                                                                  │
│   14 │                                                                                           │
│   15 │   @functools.wraps(fn)                                                                    │
│   16 │   def inner(*args, **kwargs):                                                             │
│ ❱ 17 │   │   return fn(*args, **kwargs)                                                          │
│   18 │                                                                                           │
│   19 │   return inner                                                                            │
│   20                                                                                             │
│                                                                                                  │
│ /opt/pytorch/torch/_dynamo/eval_frame.py:920 in result_capturing_wrapper                         │
│                                                                                                  │
│    917 │   │   │                                                                                 │
│    918 │   │   │   graph_captured_input = graph_inputs                                           │
│    919 │   │   │   assert graph is not None                                                      │
│ ❱  920 │   │   │   graph_captured_result = graph(*graph_inputs)                                  │
│    921 │   │   │   return graph_captured_result                                                  │
│    922 │   │                                                                                     │
│    923 │   │   return result_capturing_wrapper                                                   │
│                                                                                                  │
│ /opt/pytorch/torch/fx/graph_module.py:678 in call_wrapped                                        │
│                                                                                                  │
│   675 │   │   │   cls._wrapped_call = _WrappedCall(cls, cls_call)  # type: ignore[attr-defined   │
│   676 │   │                                                                                      │
│   677 │   │   def call_wrapped(self, *args, **kwargs):                                           │
│ ❱ 678 │   │   │   return self._wrapped_call(self, *args, **kwargs)                               │
│   679 │   │                                                                                      │
│   680 │   │   cls.__call__ = call_wrapped                                                        │
│   681                                                                                            │
│                                                                                                  │
│ /opt/pytorch/torch/fx/graph_module.py:284 in __call__                                            │
│                                                                                                  │
│   281 │   │   │   │   │     file=sys.stderr)                                                     │
│   282 │   │   │   │   raise e.with_traceback(None)                                               │
│   283 │   │   │   else:                                                                          │
│ ❱ 284 │   │   │   │   raise e                                                                    │
│   285                                                                                            │
│   286 @compatibility(is_backward_compatible=True)                                                │
│   287 class GraphModule(torch.nn.Module):                                                        │
│                                                                                                  │
│ /opt/pytorch/torch/fx/graph_module.py:274 in __call__                                            │
│                                                                                                  │
│   271 │   │   │   if self.cls_call is not None:                                                  │
│   272 │   │   │   │   return self.cls_call(obj, *args, **kwargs)                                 │
│   273 │   │   │   else:                                                                          │
│ ❱ 274 │   │   │   │   return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[mi   │
│   275 │   │   except Exception as e:                                                             │
│   276 │   │   │   assert e.__traceback__                                                         │
│   277 │   │   │   topmost_framesummary: traceback.FrameSummary = \                               │
│                                                                                                  │
│ /opt/pytorch/torch/nn/modules/module.py:1522 in _wrapped_call_impl                               │
│                                                                                                  │
│   1519 │   │   if self._compiled_call_impl is not None:                                          │
│   1520 │   │   │   return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]        │
│   1521 │   │   else:                                                                             │
│ ❱ 1522 │   │   │   return self._call_impl(*args, **kwargs)                                       │
│   1523 │                                                                                         │
│   1524 │   def _call_impl(self, *args, **kwargs):                                                │
│   1525 │   │   forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.fo  │
│                                                                                                  │
│ /opt/pytorch/torch/nn/modules/module.py:1531 in _call_impl                                       │
│                                                                                                  │
│   1528 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1529 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1530 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1531 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1532 │   │                                                                                     │
│   1533 │   │   try:                                                                              │
│   1534 │   │   │   result = None                                                                 │
│ <eval_with_key>.0:16 in forward                                                                  │
│                                                                                                  │
│ /opt/pytorch/torch/nn/modules/module.py:1522 in _wrapped_call_impl                               │
│                                                                                                  │
│   1519 │   │   if self._compiled_call_impl is not None:                                          │
│   1520 │   │   │   return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]        │
│   1521 │   │   else:                                                                             │
│ ❱ 1522 │   │   │   return self._call_impl(*args, **kwargs)                                       │
│   1523 │                                                                                         │
│   1524 │   def _call_impl(self, *args, **kwargs):                                                │
│   1525 │   │   forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.fo  │
│                                                                                                  │
│ /opt/pytorch/torch/nn/modules/module.py:1531 in _call_impl                                       │
│                                                                                                  │
│   1528 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1529 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1530 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1531 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1532 │   │                                                                                     │
│   1533 │   │   try:                                                                              │
│   1534 │   │   │   result = None                                                                 │
│                                                                                                  │
│ /opt/pytorch/torch/nn/modules/sparse.py:162 in forward                                           │
│                                                                                                  │
│   159 │   │   │   │   self.weight[self.padding_idx].fill_(0)                                     │
│   160 │                                                                                          │
│   161 │   def forward(self, input: Tensor) -> Tensor:                                            │
│ ❱ 162 │   │   return F.embedding(                                                                │
│   163 │   │   │   input, self.weight, self.padding_idx, self.max_norm,                           │
│   164 │   │   │   self.norm_type, self.scale_grad_by_freq, self.sparse)                          │
│   165                                                                                            │
│                                                                                                  │
│ /opt/pytorch/torch/nn/functional.py:2238 in embedding                                            │
│                                                                                                  │
│   2235 │   │   #   torch.embedding_renorm_                                                       │
│   2236 │   │   # remove once script supports set_grad_enabled                                    │
│   2237 │   │   _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)                    │
│ ❱ 2238 │   return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)        │
│   2239                                                                                           │
│   2240                                                                                           │
│   2241 def embedding_bag(                                                                        │
│                                                                                                  │
│ /opt/pytorch/torch/utils/_stats.py:20 in wrapper                                                 │
│                                                                                                  │
│   17 │   │   if fn.__qualname__ not in simple_call_counter:                                      │
│   18 │   │   │   simple_call_counter[fn.__qualname__] = 0                                        │
│   19 │   │   simple_call_counter[fn.__qualname__] = simple_call_counter[fn.__qualname__] + 1     │
│ ❱ 20 │   │   return fn(*args, **kwargs)                                                          │
│   21 │   return wrapper                                                                          │
│   22                                                                                             │
│                                                                                                  │
│ /opt/pytorch/torch/_subclasses/fake_tensor.py:1040 in __torch_dispatch__                         │
│                                                                                                  │
│   1037 │   │   │   return NotImplemented                                                         │
│   1038 │   │                                                                                     │
│   1039 │   │   with fake_mode:  # type: ignore[attr-defined]                                     │
│ ❱ 1040 │   │   │   return func(*args, **kwargs)                                                  │
│   1041 │                                                                                         │
│   1042 │   @staticmethod                                                                         │
│   1043 │   def _find_common_device(func, args, kwargs) -> Tuple[torch.device, bool]:             │
│                                                                                                  │
│ /opt/pytorch/torch/_ops.py:437 in __call__                                                       │
│                                                                                                  │
│   434 │   │   )                                                                                  │
│   435 │                                                                                          │
│   436 │   def __call__(self, *args, **kwargs):                                                   │
│ ❱ 437 │   │   return self._op(*args, **kwargs or {})                                             │
│   438 │                                                                                          │
│   439 │   def __hash__(self):                                                                    │
│   440 │   │   return hash(self._op)                                                              │
│                                                                                                  │
│ /opt/pytorch/torch/utils/_stats.py:20 in wrapper                                                 │
│                                                                                                  │
│   17 │   │   if fn.__qualname__ not in simple_call_counter:                                      │
│   18 │   │   │   simple_call_counter[fn.__qualname__] = 0                                        │
│   19 │   │   simple_call_counter[fn.__qualname__] = simple_call_counter[fn.__qualname__] + 1     │
│ ❱ 20 │   │   return fn(*args, **kwargs)                                                          │
│   21 │   return wrapper                                                                          │
│   22                                                                                             │
│                                                                                                  │
│ /opt/pytorch/torch/_subclasses/fake_tensor.py:1160 in __torch_dispatch__                         │
│                                                                                                  │
│   1157 │   def __torch_dispatch__(self, func, types, args=(), kwargs=None):                      │
│   1158 │   │   assert self not in _get_current_dispatch_mode_stack(), func                       │
│   1159 │   │   try:                                                                              │
│ ❱ 1160 │   │   │   return self.dispatch(func, types, args, kwargs)                               │
│   1161 │   │   except TypeError:                                                                 │
│   1162 │   │   │   log.exception("fake tensor raised TypeError")                                 │
│   1163 │   │   │   raise                                                                         │
│                                                                                                  │
│ /opt/pytorch/torch/_subclasses/fake_tensor.py:1271 in dispatch                                   │
│                                                                                                  │
│   1268 │   │   │   args,                                                                         │
│   1269 │   │   │   kwargs,                                                                       │
│   1270 │   │   │   flat_arg_fake_tensors,                                                        │
│ ❱ 1271 │   │   ) = self.validate_and_convert_non_fake_tensors(func, converter, args, kwargs)     │
│   1272 │   │                                                                                     │
│   1273 │   │   # The current constant handling only support tracing systems                      │
│   1274 │   │   # (aot autograd, torchdynamo) where each operation is run consecutively.          │
│                                                                                                  │
│ /opt/pytorch/torch/_subclasses/fake_tensor.py:1467 in validate_and_convert_non_fake_tensors      │
│                                                                                                  │
│   1464 │   │   │   flat_arg_fake_tensors.append(x)                                               │
│   1465 │   │   │   return x                                                                      │
│   1466 │   │                                                                                     │
│ ❱ 1467 │   │   args, kwargs = tree_map_only(                                                     │
│   1468 │   │   │   torch.Tensor,                                                                 │
│   1469 │   │   │   validate,                                                                     │
│   1470 │   │   │   (args, kwargs),                                                               │
│                                                                                                  │
│ /opt/pytorch/torch/utils/_pytree.py:393 in tree_map_only                                         │
│                                                                                                  │
│   390 │   ...                                                                                    │
│   391                                                                                            │
│   392 def tree_map_only(ty: TypeAny, fn: FnAny[Any], pytree: PyTree) -> PyTree:                  │
│ ❱ 393 │   return tree_map(map_only(ty)(fn), pytree)                                              │
│   394                                                                                            │
│   395 def tree_all(pred: Callable[[Any], bool], pytree: PyTree) -> bool:                         │
│   396 │   flat_args, _ = tree_flatten(pytree)                                                    │
│                                                                                                  │
│ /opt/pytorch/torch/utils/_pytree.py:323 in tree_map                                              │
│                                                                                                  │
│   320                                                                                            │
│   321 def tree_map(fn: Any, pytree: PyTree) -> PyTree:                                           │
│   322 │   flat_args, spec = tree_flatten(pytree)                                                 │
│ ❱ 323 │   return tree_unflatten([fn(i) for i in flat_args], spec)                                │
│   324                                                                                            │
│   325 Type2 = Tuple[Type[T], Type[S]]                                                            │
│   326 Type3 = Tuple[Type[T], Type[S], Type[U]]                                                   │
│                                                                                                  │
│ /opt/pytorch/torch/utils/_pytree.py:323 in <listcomp>                                            │
│                                                                                                  │
│   320                                                                                            │
│   321 def tree_map(fn: Any, pytree: PyTree) -> PyTree:                                           │
│   322 │   flat_args, spec = tree_flatten(pytree)                                                 │
│ ❱ 323 │   return tree_unflatten([fn(i) for i in flat_args], spec)                                │
│   324                                                                                            │
│   325 Type2 = Tuple[Type[T], Type[S]]                                                            │
│   326 Type3 = Tuple[Type[T], Type[S], Type[U]]                                                   │
│                                                                                                  │
│ /opt/pytorch/torch/utils/_pytree.py:374 in inner                                                 │
│                                                                                                  │
│   371 │   │   @functools.wraps(f)                                                                │
│   372 │   │   def inner(x: T) -> Any:                                                            │
│   373 │   │   │   if isinstance(x, ty):                                                          │
│ ❱ 374 │   │   │   │   return f(x)                                                                │
│   375 │   │   │   else:                                                                          │
│   376 │   │   │   │   return x                                                                   │
│   377 │   │   return inner                                                                       │
│                                                                                                  │
│ /opt/pytorch/torch/_subclasses/fake_tensor.py:1455 in validate                                   │
│                                                                                                  │
│   1452 │   │   │   │   │   )                                                                     │
│   1453 │   │   │   │   if not self.allow_non_fake_inputs:                                        │
│   1454 │   │   │   │   │   # import pdb; pdb.set_trace()                                         │
│ ❱ 1455 │   │   │   │   │   raise Exception(                                                      │
│   1456 │   │   │   │   │   │   f"Please convert all Tensors to FakeTensors first or instantiate  │
│   1457 │   │   │   │   │   │   f"with 'allow_non_fake_inputs'. Found in {render_call(func, args  │
│   1458 │   │   │   │   │   )                                                                     │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
Exception: Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in 
aten.embedding.default(Parameter(FakeTensor(..., size=(50257, 2), requires_grad=True)), tensor([...], size=(1, 3)))
```

### Versions

Latest pytorch main, but the complete list is below: 

```
Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] flake8-bugbear==23.3.23
[pip3] flake8-comprehensions==3.12.0
[pip3] flake8-executable==2.1.3
[pip3] flake8-logging-format==0.9.0
[pip3] flake8-pyi==23.3.1
[pip3] flake8-simplify==0.19.3
[pip3] mypy==0.960
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.3
[pip3] pytorch-triton==2.1.0+7d1a95b046
[pip3] torch==2.1.0a0+gita396dc5
[pip3] torch-nebula==0.16.2
[pip3] torch-ort==1.16.0.dev20230626
[pip3] torchaudio==2.1.0a0+7096829
[pip3] torchdata==0.7.0a0+901b483
[pip3] torchmetrics==1.0.0rc0
[pip3] torchsnapshot==0.1.0
[pip3] torchtext==0.16.0a0+60bea66
[pip3] torchvision==0.16.0a0+52eb503
[conda] magma-cuda117             2.6.1                         1    pytorch
[conda] mkl                       2023.1.0         h6d00ec8_46342  
[conda] mkl-include               2023.1.0         h06a4308_46342  
[conda] numpy                     1.24.3                   pypi_0    pypi
[conda] pytorch-triton            2.1.0+7d1a95b046          pypi_0    pypi
[conda] torch                     2.1.0a0+gita396dc5           dev_0    <develop>
[conda] torch-nebula              0.16.2                   pypi_0    pypi
[conda] torch-ort                 1.16.0.dev20230626           dev_0    <develop>
[conda] torchaudio                2.1.0a0+7096829          pypi_0    pypi
[conda] torchdata                 0.7.0a0+901b483          pypi_0    pypi
[conda] torchmetrics              1.0.0rc0                 pypi_0    pypi
[conda] torchsnapshot             0.1.0                    pypi_0    pypi
[conda] torchtext                 0.16.0a0+60bea66          pypi_0    pypi
[conda] torchvision               0.16.0a0+52eb503          pypi_0    pypi
```

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @aakhundov