Iterative magnitude pruning: weird results with torch.compile

sree_harsha: <p>Hello!</p>
<p>model = resnet18 with 10 classes, first layer + maxpool for cifar10 as usual and last linear is now a conv1d. For masking/sparsity purposes a mask is stored as a buffer and applied during the forward pass.</p>
<ol>
<li>I use amp training (bfloat16) + torch.compile with max-autotune and reduce-overhead.</li>
<li>I train the dense model and get test accuracy say = 93.6% and then evaluate the checkpoint separately and its consistent = 93.6% ± 0.02</li>
<li>When I train my first “sparse” model after pruning 20%, before which I save the “_orig_mod” as the state_dict</li>
<li>reload the weights, prune and then compile and retrain</li>
<li>When I evaluate the checkpoints in a separate script with the exact same test-loader – it just collapses. Not sure what I’m doing wrong with compile</li>
</ol>
<p>Left is with torch.compile and the right is without torch.compile – and that’s the <em>only</em> difference.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/e/8e8c90f3d595538adfc611443dd011c77e42de8a.png" data-download-href="https://discuss.pytorch.org/uploads/default/8e8c90f3d595538adfc611443dd011c77e42de8a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/e/8e8c90f3d595538adfc611443dd011c77e42de8a_2_690x344.png" alt="image" data-base62-sha1="kl2XG0PhVZQ74Mo1yudR14cxBFg" width="690" height="344" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/e/8e8c90f3d595538adfc611443dd011c77e42de8a_2_690x344.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/e/8e8c90f3d595538adfc611443dd011c77e42de8a_2_1035x516.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/e/8e8c90f3d595538adfc611443dd011c77e42de8a_2_1380x688.png 2x" data-dominant-color="FCFBFC"><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">3990×1990 343 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>EDIT:<br>
Fixed it partially by adding:</p>
<pre><code class="lang-auto">torch._dynamo.config.guard_nn_modules=True
</code></pre>
<p>to enable recompile after mask changes – but there’s still some discrepancy.</p>

sree_harsha: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> any thoughts here?</p>

ptrblck: <p>Just to make sure I understand the use case: the pruning strategy works fine when <code>torch.compile</code> is not used and you were able to reduce the accuracy differences via <code>torch._dynamo.config.guard_nn_modules=True</code>?<br>
CC <a class="mention" href="/u/marksaroufim">@marksaroufim</a> as you might have seen similar issues before.</p>

sree_harsha: <p>Yes, that is correct.</p>
<p>Blue line (in every graph)<br>
When I train a model and evaluate within the training harness</p>
<p>Red Line:<br>
When I load the checkpoint externally in a separate file and evaluate on the exact same test set.</p>
<p>w/o <code>torch.compile</code>:<br>
Both are consistent</p>
<p>w/ just <code>torch.compile</code>:<br>
The accuracy is just way off, my sense is that the buffers are not getting updated (i.e. the masks) in the compiled model.</p>
<p>w/ <code>torch.compile</code> and <code>torch._dynamo.config.guard_nn_modules=True</code><br>
Much better – but still some inconsistency.</p>
<p>For additional context:</p>
<ol>
<li>Using AMP with bfloat16</li>
<li>The mask is implemented as a buffer at each layer via a subclass of nn.Conv2d, applied directly to weight during the forward pass.</li>
<li>I save the model after every training ‘level’ as <code>torch.save(model._orig_mod.state_dict())</code></li>
</ol>
<p>Thanks for looking into this!</p>