I am currently working on model quantization with graph export using torch._dynamo.export and have encountered issues with torch function call, let’s say it’s torch.add.
In the exported graph, we can only obtain the call_function node of torch.add. If there is more than one torch.add in the graph, users cannot set different qconfig for each torch.add unless they manipulate the exported graph directly, which is challenging. When I use FX symbolic trace, I could wrap the torch.add in an AddModule and make that AddModule a leaf node by override the is_leaf_module method of the CustomTracer. However, I haven’t found a similar feature in dynamo.
I have observed that the modules under the torch.nn namespace are represented as call_modules in the dynamo exported graph and I have attempted to define an AddModule, place it directly in the torch.nn namespace, and I can find the AddModule’s call_module node in the graph.
I believe there might be some configurations to make dynamo keep call_module instead of tracing into the forward of Module. I have tried changing configurations like skipfiles_inline_module_allowlist in torch._dynamo.config to instruct dynamo to treat AddModule as call_module, but these changes didn’t work.
So, my question is: Is there any way to customize the leaf node when using dynamo?