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
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
torch._dynamo.config to instruct dynamo to treat
call_module, but these changes didn’t work.
So, my question is: Is there any way to customize the leaf node when using dynamo?