[dynamo][guards] Skip guards on empty nn module hooks (#138942)

anijain2305 · pytorchmergebot · commit e80fe7f13a06 · 2024-10-29T02:11:47.000Z
This brings some unsoundness in guards. Earlier we were skipping empty nn module hooks dict guard only on inbuilt nn modules, but as seen in #138386, there could be still be significant guard overhead. With this PR, we reduce the guard eval latency from 420 us to 280 us (1.5x reduction). Pull Request resolved: #138942 Approved by: https://github.com/ezyang, https://github.com/jansel ghstack dependencies: #139040, #138954
diff --git a/test/dynamo/test_activation_checkpointing.py b/test/dynamo/test_activation_checkpointing.py
@@ -519,6 +519,7 @@ def _factory_fn():
                 mod_no_hook, backend, x, fullgraph=True, compiled_autograd=True
             )
 
+        torch._dynamo.reset()
         mod_with_hook, x, backend = _factory_fn()
         mod_with_hook.submod.register_forward_hook(my_post_forward_hook)
         mod_with_hook_fwd_outputs = set()
diff --git a/test/dynamo/test_hooks.py b/test/dynamo/test_hooks.py
@@ -796,6 +796,41 @@ def forward(self, x):
 
         self.assertEqual(cnts.frame_count, 1)
 
+    @torch._dynamo.config.patch(skip_nnmodule_hook_guards=False)
+    def test_nnmodule_hook_guards(self):
+        # Compile a model and then apply a hook
+
+        class Mod(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear = torch.nn.Linear(16, 16)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        cnts = torch._dynamo.testing.CompileCounter()
+
+        mod = Mod()
+
+        def fn(x):
+            return mod(x)
+
+        opt_fn = torch.compile(fn, backend=cnts)
+
+        x = torch.ones(16, 16)
+        opt_fn(x)
+
+        # Register a hook
+        def forward_hook(self, inputs, out):
+            return out * 2
+
+        mod.register_forward_hook(forward_hook)
+
+        ref = fn(x)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+        self.assertEqual(cnts.frame_count, 2)
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/torch/_dynamo/variables/nn_module.py b/torch/_dynamo/variables/nn_module.py
@@ -1080,9 +1080,7 @@ def var_getattr(self, tx: "InstructionTranslator", name):
         ):
             # For empty hooks, make an EMPTY_NN_MODULE_HOOKS_DICT. This allows us to control the installation of empty
             # hooks guard via skip_nnmodule_hook_guards
-            if not tx.output.side_effects.has_pending_mutation_of_attr(
-                self, name
-            ) and self.value.__module__.startswith(("torch.nn.", "torch.ao.")):
+            if not tx.output.side_effects.has_pending_mutation_of_attr(self, name):
                 hooks_dict = getattr(self.value, name)
                 if isinstance(hooks_dict, dict) and len(hooks_dict) == 0:
                     if self.source:

Original file line number	Diff line number	Diff line change
`@@ -519,6 +519,7 @@ def _factory_fn():`
`519`	`519`	`mod_no_hook, backend, x, fullgraph=True, compiled_autograd=True`
`520`	`520`	`)`
`521`	`521`
	`522`	`+ torch._dynamo.reset()`
`522`	`523`	`mod_with_hook, x, backend = _factory_fn()`
`523`	`524`	`mod_with_hook.submod.register_forward_hook(my_post_forward_hook)`
`524`	`525`	`mod_with_hook_fwd_outputs = set()`