Removing Return Statement in Module Forward Causes 30+ms Backward Slowdown - Why?

xgbj: <p>Hello community members,</p>
<p>We recently encountered a puzzling performance issue during training and would appreciate your insights. The core problem is: ​<strong>​Removing a <code>return</code> statement in our module’s <code>forward()</code> method leads to a 30+ms slowdown in backward propagation, despite no logical changes to code execution.​</strong>​ Here’s the context:</p>
<pre><code class="lang-auto">def forward(self, batch_extrinsic):  
    if self._cache_bev_to_camera_projection is None:  # Cache-initialization branch  
        with torch.no_grad():  
            # ~20 tensor operations for projection/mask generation  
            bev_to_camera_projection, bev_mask = self.project_lidar_to_camera(...)  
            # Cache results  
            self._cache_bev_to_camera_projection = bev_to_camera_projection[0:1]  
            self._cache_bev_mask = bev_mask[0:1]  
        return bev_to_camera_projection, bev_mask  # Removing this line slows backward  

    # Cached branch  
    return self._cache_bev_to_camera_projection.repeat(...), self._cache_bev_mask.repeat(...)
</code></pre>
<h3><a name="p-470622-key-observations-1" class="anchor" href="#p-470622-key-observations-1"></a>Key Observations</h3>
<ol>
<li>​<strong>​With <code>return</code> statement​</strong>​:</li>
</ol>
<ul>
<li>The cache-initialization branch runs ​<strong>​only once​</strong>​ (first forward pass).</li>
<li>Backward pass time is normal.</li>
</ul>
<ol start="2">
<li>​<strong>​Without <code>return</code> statement​</strong>​:</li>
</ol>
<ul>
<li>Cache initialization still works correctly.</li>
<li>​<strong>​Backward pass slows down by 30ms​</strong>​, despite identical forward logic.</li>
</ul>
<ol start="3">
<li>​<strong>​Workaround​</strong>​:</li>
</ol>
<ul>
<li>Rewriting with <code>if-else</code> (instead of early <code>return</code>) restores normal backward performance.</li>
</ul>
<h3><a name="p-470622-hypotheses-questions-2" class="anchor" href="#p-470622-hypotheses-questions-2"></a>Hypotheses &amp; Questions</h3>
<p>We suspect this relates to ​<strong>​PyTorch’s computation graph construction​</strong>​, but need expert input:</p>
<ol>
<li>​<strong>​Gradient Tracking​</strong>​:</li>
</ol>
<ul>
<li>Despite <code>torch.no_grad()</code>, could returning the uncached tensors create unintended gradient dependencies?</li>
<li>Does <code>repeat()</code> on cached tensors retain graph connections to the initial computation?</li>
</ul>
<ol start="2">
<li>​<strong>​Graph Optimization​</strong>​:</li>
</ol>
<ul>
<li>Does the early <code>return</code> help PyTorch prune unused graph branches during backward?</li>
<li>Why does the <code>if-else</code> structure resolve the slowdown?</li>
</ul>
<ol start="3">
<li>​<strong>​Autograd Behavior​</strong>​:</li>
</ol>
<ul>
<li>How does conditional branching affect graph retention/rebuilding overhead?</li>
</ul>
<p>Any insights into why removing a logically redundant <code>return</code> impacts backward performance would be greatly appreciated!</p>
<p>Thank you!</p>

ptrblck: <p>Could you post a minimal and executable code snippet reproducing this issue, please?</p>