Dynamo Graph Capture automatic transform from `get_attr` to `placeholder`?

username: wmhst7
cooked: <p>I’m exploring the differences in graph capturing behavior between using torchdynamo and torch.fx symbolic_trace. Specifically, I’ve noticed that when tracing models using Dynamo, the <code>get_attr</code> nodes often get automatically converted to <code>placeholder</code> nodes.</p>
<p>Here’s a simple runnable demo that illustrates what I’m observing:</p>
<pre><code class="lang-auto">import torch
from torch.fx import GraphModule, Tracer, Graph
from torch.export import export

class CustomModel(torch.nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.w1 = torch.nn.Parameter(torch.empty(hidden_size, hidden_size))

    def forward(self, x):
        x = torch.mm(x, self.w1)
        # x = gelu(x)  # Uncomment for non-linear operations
        x = x.sum()
        return (x,)

if __name__ == "__main__":
    torch.set_default_dtype(torch.bfloat16)
    with torch.device("meta"):
        hidden_size = 1024
        model = CustomModel(hidden_size)
        inp = torch.zeros(2, hidden_size, requires_grad=True)
        
        tracer = Tracer()
        graph = tracer.trace(model)
        graph.print_tabular()
        
        exported_program: torch.export.ExportedProgram = export(model, args=(inp,))
        gm = exported_program.graph_module
        gm.graph.print_tabular()

</code></pre>
<p>The output using <code>torch.fx</code> directly vs. using <code>export</code> (which I presume uses Dynamo internally) shows different behaviors. In the FX trace, parameters are retained as <code>get_attr</code>, whereas in the Dynamo-based trace, they are converted to <code>placeholder</code>.</p>
<p>Output using FX:</p>
<pre><code class="lang-auto">opcode         name    target                                                 args         kwargs
-------------  ------  -----------------------------------------------------  -----------  --------
placeholder    x       x                                                      ()           {}
get_attr       w1      w1                                                     ()           {}
call_function  mm      &lt;built-in method mm of type object at 0x7f661c873500&gt;  (x, w1)      {}
call_method    sum_1   sum                                                    (mm,)        {}
output         output  output                                                 ((sum_1,),)  {}
</code></pre>
<p>Output using Dynamo:</p>
<pre><code class="lang-auto">opcode         name    target            args         kwargs
-------------  ------  ----------------  -----------  --------
placeholder    p_w1    p_w1              ()           {}
placeholder    x       x                 ()           {}
call_function  mm      aten.mm.default   (x, p_w1)    {}
call_function  sum_1   aten.sum.default  (mm,)        {}
output         output  output            ((sum_1,),)  {}
</code></pre>
<p>I am curious why this happens, and if there’s a way to control or prevent this behavior when using Dynamo. Any insights or recommendations on how to handle this discrepancy would be greatly appreciated.</p>

username: wmhst7
cooked: <p>Update:<br>
It seems it is dynamo’s feature. And for the meta(shape/dytpe…) information, dynamo graph saves them in <code>node.meta</code>.</p>

username: wmhst7
cooked: <p>Update:</p>
<p>Dynamo will automatically functionalize the graph, meaning that all input parameters and buffers are treated as graph inputs, and the entire graph is seen as a large forward function.</p>
<p>If you use <code>export</code> to get an <code>ExportedProgram</code> type, you can then call <code>torch.export.unflatten(exported)</code> to transform the graph back into an <code>UnflattenedModule</code>. In this case, the graph structure will have different hierarchical levels, such as <code>getattr</code> and <code>call_module</code>.</p>
<p>However, if you want to capture both the <strong>forward and backward</strong> (joint) graphs, you would use <code>aot_export_module</code>, which returns a <code>torch.fx.GraphModule</code> and a <code>GraphSignature</code>. This approach uses Dynamo and AOTAutoGrad to capture the graph at the lower level.</p>
<p>At this point, if we want to get an <code>ExportedProgram</code> and use <code>unflatten</code>, it is not feasible because the <code>TreeSpec</code> information in the <code>GraphSignature</code> is actually empty.</p>
<p>Nonetheless, the <code>GraphSignature</code> contains information like <code>inputs_to_parameters</code>, so we can still manually obtain the source of the placeholders, but we cannot rebuild the submodule structure.</p>