How to manipulate the input of `output` node of FX Graph generated by torchdynamo?

username: xTayEx
cooked: <p>Hello everyone! I’m currently using TorchDynamo to generate the FX graph of a <a href="https://huggingface.co/bert-base-uncased" rel="noopener nofollow ugc">bert-base-uncased</a> model. Specifically, I write a custom compile backend which only print the FX graph. I also employ the decomposition rule from inductor in order to decompose large operators to aten/prims operators. Here is the code I use</p>
<pre data-code-wrap="python"><code class="lang-python">from transformers import BertTokenizer, BertModel
import torch._dynamo as dynamo
from torch._inductor.decomposition import decompositions as inductor_decomp
from torch._functorch.aot_autograd import aot_module_simplified


def custom_backend(gm, inputs):

    def _compiler(_gm, _):
        _gm.graph.print_tabular()
        return _gm.forward

    return aot_module_simplified(gm,
                                 inputs,
                                 fw_compiler=_compiler,
                                 decompositions=inductor_decomp)


tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
text = "Hello world!"
encoded_text = tokenizer(text, return_tensors="pt")

model = BertModel.from_pretrained("bert-base-uncased")
model.eval()
model_compiled = dynamo.optimize(backend=custom_backend)(model)

model_compiled(**encoded_text)
</code></pre>
<p>However, the FX graph output by the above code is a little less than satisfactory. The input list of the <code>output</code> node is so long. Most of the input values are unnecessary for my case. Here is the <code>output</code> node printed <code>torch.fx.graph.Graph</code>:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd70eaad4cd7d28cbe892c1244f5a48528fb55e6.png" data-download-href="https://discuss.pytorch.org/uploads/default/bd70eaad4cd7d28cbe892c1244f5a48528fb55e6" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd70eaad4cd7d28cbe892c1244f5a48528fb55e6.png" alt="image" data-base62-sha1="r1SadEPRlOJxLxAGpfPaDHct9eC" width="690" height="128" data-dominant-color="393D53"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1920×357 48.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>So my question is, what causes such a long input list of <code>output</code> node? Is there any method to manipulate the FX graph capturing process of TorchDynamo so that the <code>output</code> node only contains input values I need? Any insight would be greatly appreciated! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: bdhirsh
cooked: <p>If you only care about inference, you should run your compiled model under no_grad:</p>
<p>with torch.no_grad():<br>
model_compiled(**encoded_text)</p>
<p>Most of those outputs in the forward graph correspond to “activations”, which torch.compile will then save for backward.</p>
<p>Since you are not running under no_grad, torch.compile (pessimistically) assumes that you might run .backward later, and so it needs to make sure when you create a forward graph that any potential activations are returned in the the forward, in case the user uses autograd.</p>

username: xTayEx
cooked: <p>Thank you for your reply! <code>torch.no_grad</code> works like a charm!</p>