[Bug] (suggested fix) nn.Parameter are not added to root after being traced in mmrazor.models.task_utils.tracer.fx.custom_tracer.build_graphmodule()

[elisa-aleman]

Describe the bug

Trying to do QAT training using OpenVINOQuantizer and CustomTracer, I ran into an issue whenever there is a torch.nn.Parameter traced in the graph.

Apologies but I need to redact the traceback heavily, however i will keep the necessary parts.

...
  File "..../site-packages/mmrazor/models/quantizers/native_quantizer.py", line 232, in prepare
    graph_module = build_graphmodule(model, traced_graph)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "..../site-packages/mmrazor/models/task_modules/tracer/fx/custom_tracer.py", line 214, in build_graphmodule
    return GraphModule(modules, fx_graph, name)
  File "..../site-packages/torch/fx/graph_module.py", line 355, in __init__
    raise RuntimeError(...
RuntimeError: Node ____ referenced target ____ but that target was not provided in ``root``!

Post related information

I patched this locally by finding what happened to the graph after the model got traced. I found this couple of lines in mmrazor.models.task_modules.tracer.fx.custom_tracer._prepare_module_dict()##L139--L144. I added comments with my conclusions.

    for nose in fx_graph.nodes:
        if node.op == 'get_attr': # nn.Parameters get this node.op as well.
            attr = _get_attrs(model, node.target)
            if isinstance(attr, nn.Module): # here nn.Parameters get ignored
                module_dict[node.target] = nn.Module()
                special_nodes.append(node)

[elisa-aleman]

I figured out a fix:

@@ -139,7 +139,10 @@ def _prepare_module_dict(model: torch.nn.Module, fx_graph):
     for nose in fx_graph.nodes:
         if node.op == 'get_attr':
             attr = _get_attrs(model, node.target)
             if isinstance(attr, nn.Module):
                 module_dict[node.target] = nn.Module()
                 special_nodes.append(node)
+            else:
+                module_dict[node.target] = attr
+                special_nodes.append(node)

patching this makes the graph build correctly if a model has nn.Parameter modules.

[elisa-aleman]

The previous fix trains without error but I have had issues deploying. This might be a better fix:

@@ -139,7 +139,10 @@ def _prepare_module_dict(model: torch.nn.Module, fx_graph):
     for nose in fx_graph.nodes:
         if node.op == 'get_attr':
             attr = _get_attrs(model, node.target)
             if isinstance(attr, nn.Module):
                 module_dict[node.target] = nn.Module()
                 special_nodes.append(node)
+            elif isinstance(attr, nn.Parameter):
+                module_dict[node.target] = attr
+                special_nodes.append(node)