[compiled autograd] add tests for bwd hooks relative firing order (pytorch#139004)

Pull Request resolved: pytorch#139004
Approved by: https://github.com/yf225
ghstack dependencies: pytorch#139003

Author: xmfan

test/inductor/test_compiled_autograd.py

@@ -2802,6 +2802,63 @@ def fn(x):
         with torch._dynamo.compiled_autograd.enable(torch.compile):
             out.backward()
 
+    @skipIfWindows(msg="node name demangling inconsistent on windows")
+    def test_backward_hook_relative_ordering_partial(self):
+        # test backward hooks for cases that CA matches eager
+
+        def fn():
+            order = []
+
+            class MyModule(nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.linear = torch.nn.Linear(10, 10, bias=False)
+
+                def forward(self, x):
+                    return self.linear(x)
+
+            x = torch.randn(10, 10)
+            module = MyModule()
+
+            def make_pre_hook(id):
+                return lambda _: order.append(f"pre_hook_{id}")
+
+            def make_post_hook(id):
+                return lambda _1, _2: order.append(f"post_hook_{id}")
+
+            count = 0
+
+            def register_hooks_on_all_nodes(nodes):
+                nonlocal count
+                for node, _ in nodes:
+                    if node is None:
+                        continue
+                    count += 1
+                    id = f"{node.name()}_{count}"
+                    node.register_prehook(make_pre_hook(id))
+                    node.register_hook(make_post_hook(id))
+                    register_hooks_on_all_nodes(node.next_functions)
+
+            loss = module(x).sum()
+            register_hooks_on_all_nodes(((loss.grad_fn, None),))
+
+            def make_tensor_pre_hook(id):
+                return lambda _: order.append(f"tensor_pre_hook_{id}")
+
+            def make_post_acc_grad_hook(id):
+                return lambda _: order.append(f"post_acc_grad_hook_{id}")
+
+            module.linear.weight.register_hook(make_tensor_pre_hook("weight"))
+
+            module.linear.weight.register_post_accumulate_grad_hook(
+                make_post_acc_grad_hook("weight")
+            )
+
+            loss.backward()
+            yield tuple(order)
+
+        self.check_output_and_recompiles(fn)
+
 
 def load_test_module(name):
     testdir = Path(__file__).absolute().parent.parent
@@ -2993,6 +3050,7 @@ def wrap_test_class(orig_cls):
     # Category: Divergence from eager
     "test_invalid_gradients",  # can't give autograd error due to inaccurate output metadata of lifted backward
     "test_autograd_node_isinstance",  # backward ctx is a fake cls and not directly a Node instance
+    "test_backward_hook_relative_ordering",  # compiled autograd collects breadth first, and module backward hook not supported
     # Uncategorized
 }
 

test/test_autograd.py

@@ -74,6 +74,7 @@
     skipIfMps,
     skipIfNoLapack,
     skipIfTorchDynamo,
+    skipIfWindows,
     slowTest,
     TestCase,
     xfailIfTorchDynamo,
@@ -4592,6 +4593,96 @@ def hook(t_):
         ):
             t.backward()
 
+    @skipIfWindows(msg="node name demangling inconsistent on windows")
+    def test_backward_hook_relative_ordering(self):
+        order = []
+
+        class MyModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(10, 10)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        x = torch.randn(10, 10, requires_grad=True)
+        module = MyModule()
+        module.register_full_backward_hook(
+            lambda _1, _2, _3: order.append(
+                "module_full_backward_hook_BackwardHookFunctionBackward0"
+            )
+        )
+
+        def make_pre_hook(id):
+            return lambda _: order.append(f"pre_hook_{id}")
+
+        def make_post_hook(id):
+            return lambda _1, _2: order.append(f"post_hook_{id}")
+
+        count = 0
+
+        def register_hooks_on_all_nodes(nodes):
+            nonlocal count
+            for node, _ in nodes:
+                count += 1
+                id = f"{node.name()}_{count}"
+                node.register_prehook(make_pre_hook(id))
+                node.register_hook(make_post_hook(id))
+                register_hooks_on_all_nodes(node.next_functions)
+
+        loss = module(x).sum()
+        register_hooks_on_all_nodes(((loss.grad_fn, None),))
+
+        def make_tensor_pre_hook(id):
+            return lambda _: order.append(f"tensor_pre_hook_{id}")
+
+        def make_post_acc_grad_hook(id):
+            return lambda _: order.append(f"post_acc_grad_hook_{id}")
+
+        x.register_hook(make_tensor_pre_hook("x"))
+        module.linear.weight.register_hook(make_tensor_pre_hook("weight"))
+        module.linear.bias.register_hook(make_tensor_pre_hook("bias"))
+
+        x.register_post_accumulate_grad_hook(make_post_acc_grad_hook("x"))
+        module.linear.weight.register_post_accumulate_grad_hook(
+            make_post_acc_grad_hook("weight")
+        )
+        module.linear.bias.register_post_accumulate_grad_hook(
+            make_post_acc_grad_hook("bias")
+        )
+
+        loss.backward()
+
+        expected_order = [
+            "pre_hook_SumBackward0_1",
+            "post_hook_SumBackward0_1",
+            "pre_hook_BackwardHookFunctionBackward_2",
+            "post_hook_BackwardHookFunctionBackward_2",
+            "pre_hook_AddmmBackward0_3",
+            "post_hook_AddmmBackward0_3",
+            "tensor_pre_hook_bias",
+            "pre_hook_torch::autograd::AccumulateGrad_4",
+            "post_acc_grad_hook_bias",
+            "post_hook_torch::autograd::AccumulateGrad_4",
+            "pre_hook_TBackward0_7",
+            "post_hook_TBackward0_7",
+            "tensor_pre_hook_weight",
+            "pre_hook_torch::autograd::AccumulateGrad_8",
+            "post_acc_grad_hook_weight",
+            "post_hook_torch::autograd::AccumulateGrad_8",
+            "pre_hook_BackwardHookFunctionBackward_5",
+            "module_full_backward_hook_BackwardHookFunctionBackward0",
+            "post_hook_BackwardHookFunctionBackward_5",
+            "tensor_pre_hook_x",
+            "pre_hook_torch::autograd::AccumulateGrad_6",
+            "post_acc_grad_hook_x",
+            "post_hook_torch::autograd::AccumulateGrad_6",
+        ]
+
+        self.assertEqual(len(expected_order), len(order))
+        for expected, actual in zip(expected_order, order):
+            self.assertEqual(expected, actual)
+
     def test_view_replay_enabled(self):
         def f(x):
             out = x.clone().view(-1)