<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">Our autodiff infrastructure have bugs on double backward for early expiration of grad_accumulator for some formulas (i.e. layer_norm, linear).</p>
<p dir="auto">It gives error like below</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="  File &quot;test/test_jit.py&quot;, line 659, in checkTrace
    grads2_ge = torch.autograd.grad(l2_ge, flattened_recording_inputs, allow_unused=allow_unused)
  File &quot;/scratch/wanchaol/local/pytorch/torch/autograd/__init__.py&quot;, line 149, in grad
    inputs, allow_unused)
RuntimeError: No grad accumulator for a saved leaf!"><pre class="notranslate"><code class="notranslate">  File "test/test_jit.py", line 659, in checkTrace
    grads2_ge = torch.autograd.grad(l2_ge, flattened_recording_inputs, allow_unused=allow_unused)
  File "/scratch/wanchaol/local/pytorch/torch/autograd/__init__.py", line 149, in grad
    inputs, allow_unused)
RuntimeError: No grad accumulator for a saved leaf!
</code></pre></div>
<p dir="auto">This is happening on master now, the bug is kinda complicated to trigger: when we add an non-trivial AD formula, sometimes non-trivial models will throw the above error, while in pure Autograd mode it runs fine.</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">I noticed this problem appears in <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="441906213" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/20284" data-hovercard-type="pull_request" data-hovercard-url="/pytorch/pytorch/pull/20284/hovercard" href="https://github.com/pytorch/pytorch/pull/20284">#20284</a> , after I added an AD formula for <code class="notranslate">linear</code> in <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="439389008" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/20039" data-hovercard-type="pull_request" data-hovercard-url="/pytorch/pytorch/pull/20039/hovercard" href="https://github.com/pytorch/pytorch/pull/20039">#20039</a>.</p>
<ol dir="auto">
<li>git fetch origin; git checkout gh/wanchaol/5/origin</li>
<li>python setup.py develop</li>
<li>python test/test_jit.py TestEndToEndHybridFrontendModels.test_vae</li>
</ol>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">The test will test in autograd mode and in JIT Autodiff mode, the autograd mode works fine, but the JIT autodiff second derivative gives the above error, which I think it related to our autodiff infrastructure early release one of the leaf Variable</p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="rocket" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f680.png">🚀</g-emoji> Feature</h2>
<p dir="auto">Hello,</p>
<p dir="auto">This feature is related to a thread from <a href="https://discuss.pytorch.org/t/prepare-qat-on-module-removes-hooks/72137" rel="nofollow">discuss forum</a></p>
<p dir="auto">I would like to propose mechanism of saving <code class="notranslate">_forward_pre_hooks</code> and <code class="notranslate">_forward_hooks</code> while model is preparing for quantization.</p>
<h2 dir="auto">Motivation</h2>
<p dir="auto">There are cases that depend on such functionality. Here I duplicate what I said about it in the forum</p>
<blockquote>
<p dir="auto">Have a look at this example of pre forward hook run. Here’s EfficientNet implementation that assumes it can be integrated as backbone to FPN (<a href="https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/efficientnet.py#L471">https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/efficientnet.py#L471</a>)<br>
So here we can see that during forward blocks are called sequentially than we collect input features from specific layers.</p>
</blockquote>
<h2 dir="auto">Pitch</h2>
<p dir="auto">We can add several lines here <a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantize.py#L337">https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantize.py#L337</a><br>
to preserve <code class="notranslate">mod</code>'s hooks in <code class="notranslate">new_mod</code></p>
<h2 dir="auto">Additional context</h2>
<p dir="auto">Additionally it is needed to discuss the possibilities of preserving hooks while fusing modules</p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jerryzh168/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jerryzh168">@jerryzh168</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jianyuh/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jianyuh">@jianyuh</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/dzhulgakov/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/dzhulgakov">@dzhulgakov</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/raghuramank100/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/raghuramank100">@raghuramank100</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jamesr66a/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jamesr66a">@jamesr66a</a></p>