Update on "Lift non-FakeTensor restriction for compile"

ydwu4 · ydwu4 · commit c0260b402e17 · 2023-08-14T14:57:25.000-07:00
Currently, we have the assertion that dynamo won't accept FakeTensor input unless we're exporting. This PR try to remove this restriction to finish #105679. cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 aakhundov [ghstack-poisoned]
diff --git a/test/test_fake_tensor.py b/test/test_fake_tensor.py
@@ -1,7 +1,7 @@
 # Owner(s): ["module: meta tensors"]
 
 from torch.testing._internal.common_utils import (
-    TestCase, run_tests, skipIfCrossRef, skipIfRocm, skipIfTorchDynamo, parametrize,
+    TestCase, TEST_WITH_TORCHDYNAMO, run_tests, skipIfCrossRef, skipIfRocm, skipIfTorchDynamo, parametrize,
     instantiate_parametrized_tests)
 import torch
 import torch._dynamo
@@ -111,6 +111,7 @@ def test_non_parameter_grad(self):
         fake_t = mode.from_tensor(t)
         self.assertEqual(fake_t.requires_grad, t.requires_grad)
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     @unittest.skipIf(not RUN_CUDA, "requires cuda")
     def test_index_cuda_with_cpu(self):
         with FakeTensorMode():
@@ -253,6 +254,7 @@ def test_fake_mode_error(self):
             with FakeTensorMode():
                 y = x[0]
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     def test_fake_grad_copy(self):
         x = torch.rand([4, 4], requires_grad=True)
         x.grad = torch.rand([4, 4])
@@ -282,6 +284,7 @@ def test_binary_op_type_promotion(self):
             self.assertEqual(out.dtype, torch.float)
             self.assertEqual(out.device.type, "cpu")
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     def test_from_numpy(self):
         with FakeTensorMode():
             x = torch.tensor(np.zeros([4, 4]))
@@ -353,6 +356,7 @@ def test_out_multi_device(self):
                 x.add_(y)
 
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     @unittest.skipIf(not RUN_CUDA, "requires cuda")
     def test_normalize_device(self):
         with FakeTensorMode():
@@ -369,6 +373,7 @@ def test_recursive_invocation(self):
             y = x + x
             self.assertTrue(mode.in_kernel_invocation)
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     @skipIfRocm
     @parametrize("allow_fallback_kernels", [False, True],
                  lambda a: 'with_fallback' if a else 'without_fallback')
@@ -540,6 +545,7 @@ def __init__(self):
         self.assertIs(mod_copied.a, mod_copied.b)
         self.assertEqual(mod_copied.b.storage()._cdata, mod_copied.a.storage()._cdata)
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     @unittest.skipIf(not RUN_CUDA, "requires cuda")
     def test_new(self):
         with FakeTensorMode():
@@ -550,13 +556,15 @@ def test_new(self):
             self.checkType(b.new(device='cuda'), "cuda", [0])
             self.checkType(a.new(torch.rand([1])), "cpu", [1])
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     def test_scalar_inputs(self):
         with FakeTensorMode():
             self.checkType(torch.div(3, 2), "cpu", [])
             ten = torch.zeros(2, dtype=torch.int32) * 2.0
             self.assertEqual(ten.dtype, torch.float)
             self.checkType(ten, "cpu", [2])
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     def test_allow_meta(self):
         def run_meta():
             with FakeTensorMode():
@@ -585,6 +593,7 @@ def f():
             self.assertEqual(r.size(), f.size())
             self.assertEqual(r.device, f.device)
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     def test_mixed_real_and_fake_inputs(self):
         class _TestPattern(torch.nn.Module):
             def __init__(self):
@@ -613,6 +622,7 @@ def forward(self, input):
             out = mod(torch.randn(1, 1, 3, 3))
         self.checkType(out, "cpu", (1, 1, 3, 3))
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     @unittest.skipIf(not RUN_CUDA, "requires cuda")
     def test_aten_copy_multi_device(self):
         with FakeTensorMode():
@@ -626,6 +636,7 @@ def test_aten_copy_multi_device(self):
         self.checkType(copy2, "cuda", (4,))
         self.checkType(out, "cpu", (4,))
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     @unittest.skipIf(not RUN_CUDA, "requires cuda")
     def test_aten_index_multi_device(self):
         with FakeTensorMode():
@@ -647,6 +658,7 @@ def test_aten_index_multi_device(self):
         self.checkType(r3, "cpu", (4, 4))
         self.checkType(r4, "cuda", (4, 4))
 
+    @unittest.skipIf(TEST_WITH_TORCHDYNAMO, "isinstance check for FakeTensor won't work with compile")
     @unittest.skipIf(not RUN_CUDA, "requires cuda")
     def test_aten_slice_scatter_multi_device(self):
         with FakeTensorMode():
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -18,7 +18,7 @@
 from torch import SymInt
 from torch._guards import GuardSource, TracingContext
 from torch._ops import HigherOrderOperator
-from torch._subclasses.fake_tensor import FakeTensor, is_fake, is_fake_with_fake_mode
+from torch._subclasses.fake_tensor import FakeTensor, is_fake
 from torch.fx.experimental.symbolic_shapes import (
     DimConstraint,
     DimDynamic,
@@ -948,12 +948,14 @@ def wrap_tensor(self, value: torch.Tensor):
         assert "tensor_dict" not in tensor_proxy.node.meta
         tensor_proxy.node.meta["tensor_dict"] = value.__dict__.copy()
 
+        # TODO: I think the result is guaranteed to be fake with
+        # ignore_subclass changes
+        fake_tensor_value = None
         example_value = tensor_variable.proxy.node.meta["example_value"]
-        assert is_fake_with_fake_mode(
-            example_value, self.tx.fake_mode
-        ), "Expect example_value to be Fakified by tx.fake_mode."
+        if is_fake(example_value):
+            fake_tensor_value = example_value
 
-        grapharg = GraphArg(source, value, False, example_value)
+        grapharg = GraphArg(source, value, False, fake_tensor_value)
         tensor_proxy.node.meta["grapharg"] = grapharg
         self.tx.output.add_symbol_bindings(grapharg)
 
@@ -1117,9 +1119,7 @@ def wrap_unspecialized_primitive(self, value):
                     example_value = unspec_var.proxy.node.meta["example_value"]
                 if is_fake(example_value):
                     fake_tensor_value = example_value
-                    assert is_fake_with_fake_mode(
-                        fake_tensor_value, self.tx.fake_mode
-                    ), (
+                    assert fake_tensor_value.fake_mode is self.tx.fake_mode, (
                         f"fake mode ({fake_tensor_value.fake_mode}) from fake tensor metadata doesn't match mode"
                         "({self.tx.fake_mode}) from InstructionTranslator"
                     )
@@ -1274,14 +1274,15 @@ def _clone_input(value):
         example_value = _clone_input(example_value)
         proxy.node.meta["example_value"] = example_value
         specialized_props = target_cls.specialize(example_value)
-        assert is_fake_with_fake_mode(
-            example_value, tx.fake_mode
-        ), "Expect all example_value to fakified by tx.fake_mode by now."
-        # Example value need to preserve the original class type by replacing the leaves
-        # of subclasses with fake_tensor. Otherwise, isinstance() calls will produce wrong results.
-        specialized_props["class_type"] = (
-            torch.nn.Parameter if is_parameter else type(example_value)
-        )
+        # TODO: not sure about this fake mode test
+        if (
+            isinstance(example_value, torch._subclasses.fake_tensor.FakeTensor)
+            and example_value.fake_mode is tx.fake_mode
+        ):
+            # NB: This will be wrong for ignore_subclass; fix it up later!
+            specialized_props["class_type"] = (
+                torch.nn.Parameter if is_parameter else torch.Tensor
+            )
 
         specialized_props["specialized_value"] = specialized_value
 
diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py
@@ -160,8 +160,8 @@ def _is_tensor_constructor(func: OpOverload):
     )
 
 
-def is_fake(x, fake_mode=None):
-    if isinstance(x, FakeTensor) and (x.fake_mode is fake_mode or fake_mode is None):
+def is_fake(x):
+    if isinstance(x, FakeTensor):
         return True
     if is_traceable_wrapper_subclass(x):
         flattened_tensors, _ = type(x).__tensor_flatten__(x)
@@ -173,10 +173,6 @@ def is_fake(x, fake_mode=None):
     return False
 
 
-def is_fake_with_fake_mode(x, fake_mode):
-    return is_fake(x, fake_mode)
-
-
 @functools.lru_cache(None)
 def get_schema_info(func):
     return torch._C._SchemaInfo(func._schema)  # type: ignore[attr-defined]
