Model forward slows down when disabling gradients

KC314: <p>I have a model where, when I run <em>only</em> the forward pass, it is slower when I do <code>set_grad_enabled(False)</code> than when I do <code>set_grad_enabled(True)</code>. My intuition is that <code>set_grad_enabled(False)</code> should only stop saving various activations and intermediate tensors for the backward pass but otherwise run the code exactly the same.</p>
<p>However, I’ve looked at the trace in nvidia’s nsight systems and I do see that different kernels are getting executed as a result of changing this flag (nothing else in my model is using this parameter to change code paths). Specifically, my linear layers take about 10x longer when disabling gradients.</p>
<p>When I enable gradients and I check the corresponding kernel that gets chosen for a particular linear layer, I see<br>
<strong>ampere_fp16_s16816gemm_fp16_256x128_ldg8_f2f_stages_64x3_tn</strong><br>
<strong>block : &lt;&lt;&lt;128, 1, 1&gt;&gt;&gt;</strong><br>
<strong>grid: &lt;&lt;&lt;&lt;17, 263, 2&gt;&gt;&gt;</strong></p>
<p>In the no grad case, looking at the same linear layer, I see the kernel<br>
<strong>ampere_fp16_s16816gemm_fp16_64x64_sliced1x2_ldg8_f2f_stages_64x5_tn</strong><br>
<strong>block : &lt;&lt;&lt;128, 1, 1&gt;&gt;&gt;</strong><br>
<strong>grid: &lt;&lt;&lt;&lt;33, 1, 33600&gt;&gt;&gt;</strong>.</p>
<p>Firstly, I don’t understand why <code>set_enable_grad</code> context manager has any effect on which kernel gets launched. Secondly, it seems the kernel being chosen in the “no grad” case is launching way too many blocks and that’s probably why the it’s so much slower.</p>
<p>Some other info, this on an NVIDIA A100, torch==2.1.0+cu118, and I’m not using <code>torch.compile</code>. I’ve tried to reproduce this by using only a single linear layer but I don’t observe the behaviour in that case. The tensors being multiplied are (1024, 2048) from the linear weights and then (2048, 1, 33600) activations</p>

ptrblck: <p>Could you post a minimal and executable code snippet reproducing the different kernel execution?</p>

KC314: <p>I’m working on this but so far I’m only able to repro it on my entire model which is complicated and spread across a bunch of files. But in the meantime, is there any insight as to what <code>set_enable_grad</code> is actually doing? It’s doing a lot more than I thought if it’s actually changing which kernels are getting launched</p>

ptrblck: <p><code>set_enable_grad(False)</code> will disable creating the computation graph and will thus not store intermediate activations needed for the <code>backward</code> pass. It should not change any kernel selection, which is why we would need to get a reproducer to debug it further.</p>

soulitzer: <p>It is possible that grad mode changes the code path for matmuls - <a href="https://github.com/pytorch/pytorch/issues/118548#issuecomment-1916022394" class="inline-onebox" rel="noopener nofollow ugc">e10cfdd8953 regresses hf_Longformer-cuda-eager on pytorch 2024-01-16 · Issue #118548 · pytorch/pytorch · GitHub</a></p>
<p>To double check, you can run the following to see what ops are running underneath your linear:</p>
<pre><code class="lang-auto">from torch.testing._internal.logging_tensor import capture_logs_with_logging_tensor_mode

with capture_logs_with_logging_tensor_mode() as logs:
     &lt;do stuff here&gt;

print('\n'.join(logs))
</code></pre>

KC314: <p>Thanks <a class="mention" href="/u/soulitzer">@soulitzer</a>! This seems like exactly what’s happening actually</p>
<p>In the case when I enable gradients, I get</p>
<pre><code class="lang-auto">...
$507: f16[33600, 1, 2056] = torch._ops.aten.constant_pad_nd.default($320, ['0', '0', '0', '0', '0', '0'], 0.0)
$508: f16[2100, 1, 2056] = torch._ops.aten.constant_pad_nd.default($506, ['0', '0', '0', '0', '0', '0'], 0.0)
$510: f16[4112, 2056] = torch._ops.aten._to_copy.default($509, dtype=torch.float16)
$511: f16[2056, 4112] = torch._ops.aten.t.default($510)
$512: f16[2100, 2056] = torch._ops.aten.view.default($508, ['2100', '2056'])
$513: f16[2100, 4112] = torch._ops.aten.mm.default($512, $511)
..
</code></pre>
<p>So it’s doing a matmul with <code>(2100, 2056)</code> times <code>(2056, 4112)</code>.</p>
<p>When I disable gradients though, I get</p>
<pre><code class="lang-auto">...
$322: f16[33600, 1, 2056] = torch._ops.aten.constant_pad_nd.default($319, ['0', '0', '0', '0', '0', '0'], 0.0)
$323: f16[2100, 1, 2056] = torch._ops.aten.constant_pad_nd.default($321, ['0', '0', '0', '0', '0', '0'], 0.0)
$325: f16[4112, 2056] = torch._ops.aten._to_copy.default($324, dtype=torch.float16)
$326: f16[2056, 4112] = torch._ops.aten.t.default($325)
$327: f16[2100, 1, 2056] = torch._ops.aten.expand.default($323, ['2100', '1', '2056'])
$328: f16[2100, 1, 2056] = torch._ops.aten.view.default($327, ['2100', '1', '2056'])
$329: f16[2100, 2056, 4112] = torch._ops.aten.expand.default($326, ['2100', '2056', '4112'])
$330: f16[2100, 2056, 4112] = torch._ops.aten.view.default($329, ['2100', '2056', '4112'])
$331: f16[2100, 1, 4112] = torch._ops.aten.bmm.default($328, $330)
...
</code></pre>
<p>This is the same linear layer and I notice it inserts an expand so that it ends up doing a <strong>batched</strong> matmul of <code>(2100, 1, 2056)</code> with <code>(2100, 2056, 4112)</code>. The end result is that it’s doing 2100 covector-matrix multiplies of <code>(1, 2056)</code> with <code>(2056, 4112)</code> which is less efficient than the first case of <code>(2100, 2056) @ (2056, 4112)</code>.</p>
<p>Based on the linked comment, it seems like the second snippet is hitting case 3 and the first snippet is hitting case 2? Weirdly I can’t reproduce the issue with just a linear layer even when I try to ensure that the tensor isn’t contiguous.</p>
<p>This ends up happening a bunch of different places in the model so it slows down validation considerably. And sometimes that 2100 dimension is much much bigger making it even more inefficient.</p>