Compiled_autograd Causes Abnormal Model Training Behavior

smartbarbarian: <p>I’m encountering an issue with enabling <strong><code>compiled_autograd</code></strong> in PyTorch, which is causing my model’s training behavior to differ significantly.</p>
<h3><a name="p-462717-problem-description-1" class="anchor" href="#p-462717-problem-description-1"></a><strong>Problem Description</strong></h3>
<p>Under the same environment, model, and dataset, toggling the following configuration:</p>
<pre><code class="lang-auto">torch._dynamo.config.compiled_autograd = True
</code></pre>
<p>produces a noticeably different trend in <strong>loss</strong> and <strong>gradient norm</strong> compared to setting it to <code>False</code>.</p>
<h3><a name="p-462717-observations-2" class="anchor" href="#p-462717-observations-2"></a><strong>Observations</strong></h3>
<ul>
<li>When <strong><code>compiled_autograd</code></strong> is <strong>off</strong> (<code>False</code>), the loss and gradient norm trends behave as expected.</li>
<li>However, turning it <strong>on</strong> (<code>True</code>) results in abnormal trends.</li>
</ul>
<h3><a name="p-462717-visual-comparison-3" class="anchor" href="#p-462717-visual-comparison-3"></a><strong>Visual Comparison</strong></h3>
<p>I’ve attached the plots below:</p>
<ul>
<li><strong>Red</strong>: <code>compiled_autograd = False</code></li>
<li><strong>Gray</strong>: <code>compiled_autograd = True</code></li>
</ul>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/b/fb104379a46a44885c58572f36ce8cf4fee84d83.png" data-download-href="https://discuss.pytorch.org/uploads/default/fb104379a46a44885c58572f36ce8cf4fee84d83" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/b/fb104379a46a44885c58572f36ce8cf4fee84d83.png" alt="image" data-base62-sha1="zP0GKWpFJqteBbaB4VcIJm8PmYX" width="373" height="238"><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">373×238 24.6 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><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/3/63ef11c1923f834996b54e8fc2684a90d99d49e9.png" data-download-href="https://discuss.pytorch.org/uploads/default/63ef11c1923f834996b54e8fc2684a90d99d49e9" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/3/63ef11c1923f834996b54e8fc2684a90d99d49e9.png" alt="image" data-base62-sha1="eg3tlqptxPmMtTxCirjRNCU8Nkd" width="378" height="247"><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">378×247 19.6 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>
<h3><a name="p-462717-environment-details-4" class="anchor" href="#p-462717-environment-details-4"></a><strong>Environment Details</strong></h3>
<ul>
<li>PyTorch version: 2.5.1</li>
<li>CUDA version:  12.1</li>
<li>GPU: NVIDIA GeForce RTX 3090</li>
<li>Operating System: ubuntu 22.04</li>
</ul>
<p>Any insights or suggestions would be greatly appreciated!</p>

ptrblck: <p>This sounds indeed like a bug. Do you see this issue in all of your tests or models or only for a specific use case? In the latter case, could you share a minimal and executable code snippet reproducing the issue?</p>