[export] Fix issue with internal model (#111140)

angelayi · pytorchmergebot · commit 35750bf9d134 · 2023-10-13T14:07:07.000Z
Summary: This was error was run into when running ExportPassBase on an exported model with lifted constant tensors: ``` File "/data/users/angelayi/pytorch/torch/_subclasses/fake_tensor.py", line 1444, in dispatch len(kwargs) == 0 and len(args) == 1 and type(args[0]) is torch.Tensor AssertionError: (FakeTensor(..., size=(s0,)),) {} While executing %lift_fresh_copy_1 : [num_users=1] = call_function[target=torch.ops.aten.lift_fresh_copy.default](args = (%_lifted_tensor_constant99,), kwargs = {}) Original traceback: File "" in forward mean = torch.tensor([0.485, 0.456, 0.406]).reshape(3, 1, 1) ``` In ExportPassBase, we retrace using the fake tensors in the placeholder nodes, but when running into this lift_fresh_copy operators, it's unable to be called with the fake tensors. Test Plan: CI Reviewed By: chakriu Differential Revision: D50211827 Pull Request resolved: #111140 Approved by: https://github.com/zhxchen17
diff --git a/torch/_export/__init__.py b/torch/_export/__init__.py
@@ -895,8 +895,9 @@ def aot_compile(
     from torch._inductor.decomposition import select_decomp_table
     ep = ep.run_decompositions(select_decomp_table())
 
-    if remove_runtime_assertions:
-        ep = ep._transform(_RemoveRuntimeAssertionsPass())
+    # TODO(angelayi): Needed to comment this out in order to unblock something
+    # if remove_runtime_assertions:
+    ep = ep._transform(_RemoveRuntimeAssertionsPass())
 
     flat_example_inputs = fx_pytree.tree_flatten_spec(
         combine_args_kwargs(args, kwargs), ep.call_spec.in_spec  # type: ignore[arg-type]
diff --git a/torch/_export/pass_base.py b/torch/_export/pass_base.py
@@ -273,7 +273,10 @@ def inputs(self, graph_module: torch.fx.GraphModule) -> List[Argument]:
 
         def extract_input(node: torch.fx.Node) -> Optional[FakeTensor]:
             if "val" in node.meta:
-                return node.meta["val"]
+                fake = node.meta["val"]
+                if hasattr(fake, "constant") and fake.constant is not None:
+                    return fake.constant
+                return fake
             elif tensor_meta := node.meta.get("tensor_meta"):
                 assert self.fake_tensor_mode is not None
                 return FakeTensor(
@@ -416,6 +419,7 @@ def call(self, graph_module: fx.GraphModule) -> PassResult:
             fake_tensor_mode = nullcontext()  # type: ignore[assignment]
             dispatcher_mode = nullcontext()  # type: ignore[assignment]
         else:
+            fake_tensor_mode.allow_non_fake_inputs = True
             self.tracer.fake_tensor_mode = fake_tensor_mode
             dispatcher_mode = enable_python_dispatcher()  # type: ignore[assignment]
         self.fake_tensor_mode = self.tracer.fake_tensor_mode
