[dynamo][decorator] Support disable on nn modules (#124185)

anijain2305 · pytorchmergebot · commit f43351718165 · 2024-04-17T16:20:34.000Z
Fixes #123979 Pull Request resolved: #124185 Approved by: https://github.com/weifengpy, https://github.com/yoyoyocmu
diff --git a/test/dynamo/test_decorators.py b/test/dynamo/test_decorators.py
@@ -84,6 +84,56 @@ def wrapper():
         # to callsites of eval_frame.innermost_fn. A warning would also be very noisy.
         w = torch._dynamo.disable(fn=wrapper, recursive=True)
 
+    def test_disable_nn_modules_forward_hook(self):
+        class SimpleLinear(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer0 = torch.nn.Linear(4, 4)
+
+            def forward(self, inp):
+                return self.layer0(torch.sigmoid(inp))
+
+        class SimpleModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layer0 = SimpleLinear()
+                self.layer1 = torch.nn.Linear(4, 4)
+
+            def forward(self, inp):
+                z = self.layer0(torch.sin(inp))
+                return self.layer1(z)
+
+        def hook(module, args):
+            inp = args[0].sigmoid()
+            return (inp,)
+
+        model = SimpleModel()
+        model.layer0.register_forward_pre_hook(hook)
+
+        # Disable my monkeypatching
+        model.layer0 = torch._dynamo.disable(model.layer0)
+
+        cnts = torch._dynamo.testing.CompileCounterWithBackend("eager")
+        opt_model = torch.compile(model, backend=cnts)
+        opt_model(torch.randn(4))
+
+        # check for no graph break
+        self.assertEqual(cnts.frame_count, 2)
+
+        gm0 = cnts.graphs[0]
+        # Check that the first graph has sin node, and no sigmoid
+        self.assertTrue(any(node.target is torch.sin for node in gm0.graph.nodes))
+        self.assertTrue(
+            all(node.target is not torch.sigmoid for node in gm0.graph.nodes)
+        )
+
+        gm1 = cnts.graphs[1]
+        # Check that the first graph does not have sigmoid. sigmoid is used in
+        # both hook and disabled module.
+        self.assertTrue(
+            all(node.target is not torch.sigmoid for node in gm1.graph.nodes)
+        )
+
     def test_allow_in_graph(self):
         cnts = torch._dynamo.testing.CompileCounter()
 
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -137,6 +137,7 @@
     AutogradFunctionVariable,
     ComptimeVariable,
     DebuggingVariable,
+    DelayGraphBreakVariable,
     GetAttrVariable,
     GetSetDescriptorVariable,
     InspectSignatureVariable,
@@ -972,6 +973,14 @@ def wrap_module(self, value: torch.nn.Module):
         if len(value.__dict__) == 0:
             unimplemented(f"uninitialized nn.Module: {typestr(value)}")
         if istype(value, OptimizedModule):
+            # Check if the optimized module was disabled
+            if inspect.getattr_static(value.forward, "_torchdynamo_disable", False):
+                # This bytecode is mostly of kind LOAD_ATTR or LOAD_METHOD. If
+                # we graph break here, Dynamo does not know how to create
+                # continuation functions for such bytecodes. So, we delay the
+                # graph break to CALL_FUNCTION.
+                return DelayGraphBreakVariable(source=self.source)
+
             self.install_guards(GuardBuilder.TYPE_MATCH)
             self.source = AttrSource(self.source, "_orig_mod")
             return self.wrap_module(value._orig_mod)
