[pt2d] module register_pre_forward_hook and register_forward_hook triggered graph break when it's root module

[wanchaol]

🐛 Describe the bug

The issue pops when the hook registered as the "root" module, where we did not inline the hook calls into the graph.
But if the hooks not registered as a root module, the dynamo graph inlining works

import torch
import torch.nn as nn

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(18, 18, bias=False)

    def forward(self, x):
        return self.net1(x)

mod = ToyModel()
mod.register_forward_pre_hook(
    lambda mod, input: input[0] + 1
)

compiled_mod = torch.compile(mod)
compiled_mod(torch.rand(18, 18))

TORCH_LOGS="graph" python test_compile.py

outputs:

[2024-01-16 14:27:41,383] [0/0] torch._dynamo.output_graph.__graph: [DEBUG] TRACED GRAPH
[2024-01-16 14:27:41,383] [0/0] torch._dynamo.output_graph.__graph: [DEBUG]  __compiled_fn_0 <eval_with_key>.0 opcode         name        target                   args             kwargs
[2024-01-16 14:27:41,383] [0/0] torch._dynamo.output_graph.__graph: [DEBUG] -------------  ----------  -----------------------  ---------------  --------
[2024-01-16 14:27:41,383] [0/0] torch._dynamo.output_graph.__graph: [DEBUG] placeholder    l_input_0_  L_input_0_               ()               {}
[2024-01-16 14:27:41,383] [0/0] torch._dynamo.output_graph.__graph: [DEBUG] call_function  add         <built-in function add>  (l_input_0_, 1)  {}
[2024-01-16 14:27:41,383] [0/0] torch._dynamo.output_graph.__graph: [DEBUG] output         output      output                   ((add,),)        {}
[2024-01-16 14:27:41,383] [0/0] torch._dynamo.output_graph.__graph: [DEBUG] 
[2024-01-16 14:27:46,032] [1/0] torch._dynamo.output_graph.__graph: [DEBUG] TRACED GRAPH
[2024-01-16 14:27:46,032] [1/0] torch._dynamo.output_graph.__graph: [DEBUG]  __compiled_fn_1 <eval_with_key>.36 opcode       name            target          args                  kwargs
[2024-01-16 14:27:46,032] [1/0] torch._dynamo.output_graph.__graph: [DEBUG] -----------  --------------  --------------  --------------------  --------
[2024-01-16 14:27:46,032] [1/0] torch._dynamo.output_graph.__graph: [DEBUG] placeholder  l_x_            L_x_            ()                    {}
[2024-01-16 14:27:46,032] [1/0] torch._dynamo.output_graph.__graph: [DEBUG] call_module  l__self___net1  L__self___net1  (l_x_,)               {}
[2024-01-16 14:27:46,032] [1/0] torch._dynamo.output_graph.__graph: [DEBUG] output       output          output          ((l__self___net1,),)  {}

Versions

main branch

cc @H-Huang @awgu @kwen2501 @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @XilunWu @tianyu-l @yf225 @ezyang @anijain2305 @rohan-varma @msaroufim @bdhirsh @zou3519 @aakhundov

[bdhirsh]

cc @voznesenskym for hooks?

[anijain2305]

Need to look deeper but for some reason Dynamo is tracing the hook separately and not part of root module forward.

[wconstab]

i think its bc the root module is the dynamo OptimizedModule wrapper thing, and we have to specially handle hooks on that level since dynamo's entry point is it's forward or something.

[xmfan]

Seems like in eager, we call into the hook:

pytorch/torch/nn/modules/module.py

Line 1566 in a9d9077

args_result = hook(self, args)

So dynamo first traces a graph for the hook, then traces another one for the forward

[xmfan]

I don't fully understand this code, but for the hooks in non-root case, we rely on this NNModuleVariable -> UnspecializedNNModuleVariable conversion + dynamo restart:

pytorch/torch/_dynamo/variables/nn_module.py

Lines 301 to 305 in b2dddcf

	if nnmodule_has_hooks(
	mod, check_forward_hooks=True, check_backward_hooks=True
	):
	# End of fn, this bubbles up and restarts tracing.
	self.convert_to_unspecialized(tx)

I tried letting dynamo inline the entire root module's __call__ in a similar manner, but there's a lot of failures due to UnspecializedNNModuleVariable not being able to handle what NNModuleVariable could

#120845 in case anyone else wants to take a jab at it