Unexpected execution time difference for identical operations on GPU

Szymon_Antoniak: <p>Hello everyone,<br>
I am training a transformer with a custom FeedForward layer and it’s causing unexpected efficiency issues. This layer entails two heavy batched matrix multiplications, implemented with torch.bmm</p>
<p>My code is basically:<br>
<code>def my_layer(x):</code><br>
<code>   x = torch.bmm(x,tensor_1)    #call it operation1</code><br>
<code>   x = torch.relu_(x)</code><br>
<code>   x = torch.bmm(x,tensor_2)    #call it operation2</code></p>
<p>In my case, the shape of x is always <code>torch.Size([1024, 256, 512])</code>, and shape of tensor_1 and tensor_2 is <code>torch.Size([1024, 512, 512])</code>.</p>
<p>As seen above, both <code>operation1</code> and <code>operation2</code> implement the exact same type of operation with tensors that are identical in shape and type (the shape of <code>x</code> stays the same throughout this code), and yet the first one takes significantly more time than the other. The timings look as follows:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/c/ac17d0f8a13bece105e901a0dfde1a4d8ad1b07e.png" data-download-href="https://discuss.pytorch.org/uploads/default/ac17d0f8a13bece105e901a0dfde1a4d8ad1b07e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/ac17d0f8a13bece105e901a0dfde1a4d8ad1b07e_2_690x288.png" alt="image" data-base62-sha1="oyp9mNiA9x2u7Gqylv73LqN2gdo" width="690" height="288" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/ac17d0f8a13bece105e901a0dfde1a4d8ad1b07e_2_690x288.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/ac17d0f8a13bece105e901a0dfde1a4d8ad1b07e_2_1035x432.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/ac17d0f8a13bece105e901a0dfde1a4d8ad1b07e_2_1380x576.png 2x" data-dominant-color="BFC5F9"><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">1522×636 14.1 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>The difference is more pronounced with mixed precision enabled:<br>
<em>&lt;picture 2 in the reply - the site only allows 1 image per post&gt;</em></p>
<p>I then took to looking at what would happen if I followed the two operations with their repeated uses, i.e. <code>operation1</code> → <code>operation2</code> → <code>operation1</code> → <code>operation1</code> → <code>operation2</code> → <code>operation2</code><br>
(relu omitted for brevity)</p>
<p>The results are as follows:<br>
<em>&lt;picture 3 in the reply&gt;</em><br>
We can see that only the first operation execution takes more time.<br>
With mixed precision enabled, even more interesting things happen:<br>
<em>&lt;picture 4 in the reply&gt;</em><br>
As we see, the first execution is the longest for a given operation, and the further ones all take the same amount of time.</p>
<p>I have two questions:</p>
<ol>
<li>Why are execution times different for the first use of <code>operation1</code> and <code>operation2</code>?</li>
<li>Why does mixed precision cause the repeated use of the same parameters to take less time to do the calculations?</li>
</ol>
<p>My goal is to shorten the time for both <code>operation1</code> and <code>operation2</code>, by bridging the gap between them and, if possible, getting them down to the time displayed in <code>operation_1_again</code> and <code>operation_2_again</code>.</p>
<p>The way I measure time is standard: i use cuda events for start and end of measurement, followed by <code>torch.cuda.synchronise()</code></p>
<p>Any insights into inner workings of CUDA (I am using A100 80GB and PyTorch 2.0) and mixed precision intricacies would be greatly appreciated.</p>

Szymon_Antoniak: <p>pic 2<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/5/35a36479d01135069bb8580fd3edcd7893de2d02.png" data-download-href="https://discuss.pytorch.org/uploads/default/35a36479d01135069bb8580fd3edcd7893de2d02" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/5/35a36479d01135069bb8580fd3edcd7893de2d02_2_690x288.png" alt="image" data-base62-sha1="7Evn6IzXpWxnZOFiFxmnPe1egM2" width="690" height="288" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/5/35a36479d01135069bb8580fd3edcd7893de2d02_2_690x288.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/5/35a36479d01135069bb8580fd3edcd7893de2d02_2_1035x432.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/5/35a36479d01135069bb8580fd3edcd7893de2d02_2_1380x576.png 2x" data-dominant-color="C5CBF9"><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">1522×636 18.2 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>

Szymon_Antoniak: <p>pic 3<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/0/f016bb9823d853f997337eaa838ae82eb2866505.png" data-download-href="https://discuss.pytorch.org/uploads/default/f016bb9823d853f997337eaa838ae82eb2866505" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/0/f016bb9823d853f997337eaa838ae82eb2866505_2_690x288.png" alt="image" data-base62-sha1="yfVisiXjIW41yTr0Y7gZS0h4MXb" width="690" height="288" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/0/f016bb9823d853f997337eaa838ae82eb2866505_2_690x288.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/0/f016bb9823d853f997337eaa838ae82eb2866505_2_1035x432.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/0/f016bb9823d853f997337eaa838ae82eb2866505_2_1380x576.png 2x" data-dominant-color="ABB3FA"><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">1522×636 21 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>

Szymon_Antoniak: <p>pic 4<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/e/be2f0cc8dcb3fbe5c98c4886c54aa4d40313e037.png" data-download-href="https://discuss.pytorch.org/uploads/default/be2f0cc8dcb3fbe5c98c4886c54aa4d40313e037" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be2f0cc8dcb3fbe5c98c4886c54aa4d40313e037_2_690x288.png" alt="image" data-base62-sha1="r8rws0QJg6AbTin7p6b4OejENjF" width="690" height="288" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be2f0cc8dcb3fbe5c98c4886c54aa4d40313e037_2_690x288.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be2f0cc8dcb3fbe5c98c4886c54aa4d40313e037_2_1035x432.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be2f0cc8dcb3fbe5c98c4886c54aa4d40313e037_2_1380x576.png 2x" data-dominant-color="D0D6F8"><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">1522×636 26 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>

ptrblck: <p>Could you check which kernels are called in both use cases and if they are the same? Also, did you allow TF32 to be used in matmuls?</p>

Szymon_Antoniak: <p>As to allowing TF32, I did not take any explicit steps to do that.<br>
As for checking which kernel are used, I’m happy to check, but what would be the preffered way of doing that?</p>

ptrblck: <p>You can use the native profiler and just follow our tutorial or e.g. Nsight Systems.</p>

Szymon_Antoniak: <p>OK, this is the context manager I used:</p>
<pre><code class="lang-auto">        with profile(
            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
            record_shapes=True,
            profile_memory=True,
            use_cuda=True,
        ) as prof:
            with record_function("cont_moe"):
</code></pre>
<p>I used the original setup with only 1 use of each operation (if the others are needed, I’m happy to run the tests again).</p>
<ol>
<li>Mixed precision enabled:</li>
</ol>
<pre><code class="lang-auto">-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                               cont_moe         1.10%      76.000us        14.69%       1.016ms       1.016ms       0.000us         0.00%       6.143ms       6.143ms           0 b           0 b       1.50 Gb           0 b             1  
                                              aten::bmm         0.98%      68.000us        14.53%       1.005ms     251.250us       1.388ms        22.59%       9.050ms       2.263ms           0 b           0 b       2.00 Gb     256.00 Mb             4  
                                               aten::to         0.04%       3.000us        11.45%     792.000us     396.000us       0.000us         0.00%       2.600ms       1.300ms           0 b           0 b       1.00 Gb           0 b             2  
                                         aten::_to_copy         0.33%      23.000us        11.41%     789.000us     394.500us       0.000us         0.00%       2.600ms       1.300ms           0 b           0 b       1.00 Gb           0 b             2  
                                    aten::empty_strided         5.68%     393.000us         5.68%     393.000us     196.500us       0.000us         0.00%       0.000us       0.000us           0 b           0 b       1.00 Gb       1.00 Gb             2  
                                            aten::copy_         0.36%      25.000us         5.63%     389.000us     129.667us       4.437ms        72.23%       4.437ms       1.479ms           0 b           0 b           0 b           0 b             3  
                                       cudaLaunchKernel         5.52%     382.000us         5.52%     382.000us      63.667us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             6  
void at::native::unrolled_elementwise_kernel&lt;at::nat...         0.00%       0.000us         0.00%       0.000us       0.000us       2.600ms        42.32%       2.600ms       1.300ms           0 b           0 b           0 b           0 b             2  
                                            aten::clone         0.07%       5.000us         0.43%      30.000us      30.000us       0.000us         0.00%       1.837ms       1.837ms           0 b           0 b     256.00 Mb           0 b             1  
                                       aten::empty_like         0.07%       5.000us         0.13%       9.000us       9.000us       0.000us         0.00%       0.000us       0.000us           0 b           0 b     256.00 Mb           0 b             1  
                                            aten::empty         0.06%       4.000us         0.06%       4.000us       4.000us       0.000us         0.00%       0.000us       0.000us           0 b           0 b     256.00 Mb     256.00 Mb             1  
cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFla...         0.04%       3.000us         0.04%       3.000us       0.750us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  
                                            aten::relu_         0.26%      18.000us         0.52%      36.000us      36.000us       0.000us         0.00%     318.000us     318.000us           0 b           0 b           0 b           0 b             1  
                                       aten::clamp_min_         0.16%      11.000us         0.26%      18.000us      18.000us     318.000us         5.18%     318.000us     318.000us           0 b           0 b           0 b           0 b             1  
                                  cudaDeviceSynchronize        85.31%       5.899ms        85.31%       5.899ms       5.899ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             1  
void at::native::elementwise_kernel&lt;128, 4, at::nati...         0.00%       0.000us         0.00%       0.000us       0.000us       1.837ms        29.90%       1.837ms       1.837ms           0 b           0 b           0 b           0 b             1  
ampere_fp16_s16816gemm_fp16_128x128_ldg8_f2f_stages_...         0.00%       0.000us         0.00%       0.000us       0.000us       1.388ms        22.59%       1.388ms     694.000us           0 b           0 b           0 b           0 b             2  
void at::native::vectorized_elementwise_kernel&lt;4, at...         0.00%       0.000us         0.00%       0.000us       0.000us     318.000us         5.18%     318.000us     318.000us           0 b           0 b           0 b           0 b             1  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 6.915ms
Self CUDA time total: 6.143ms
</code></pre>
<ol start="2">
<li>Mixed precision disabled:</li>
</ol>
<pre><code class="lang-auto">
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                               cont_moe         0.28%      58.000us         4.28%     884.000us     884.000us       0.000us         0.00%      17.387ms      17.387ms           0 b           0 b       1.00 Gb           0 b             1  
                                              aten::bmm         1.85%     381.000us         3.85%     794.000us     397.000us      14.583ms        83.87%      16.763ms       8.382ms           0 b           0 b       1.00 Gb     512.00 Mb             2  
                                            aten::clone         0.02%       5.000us         1.93%     398.000us     398.000us       0.000us         0.00%       2.180ms       2.180ms           0 b           0 b     512.00 Mb           0 b             1  
                                       aten::empty_like         0.03%       7.000us         0.05%      11.000us      11.000us       0.000us         0.00%       0.000us       0.000us           0 b           0 b     512.00 Mb           0 b             1  
                                            aten::empty         0.02%       4.000us         0.02%       4.000us       4.000us       0.000us         0.00%       0.000us       0.000us           0 b           0 b     512.00 Mb     512.00 Mb             1  
                                            aten::copy_         0.06%      13.000us         1.85%     382.000us     382.000us       2.180ms        12.54%       2.180ms       2.180ms           0 b           0 b           0 b           0 b             1  
                                       cudaLaunchKernel         1.88%     388.000us         1.88%     388.000us      97.000us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             4  
cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFla...         0.01%       3.000us         0.01%       3.000us       0.250us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b            12  
                                            aten::relu_         0.08%      16.000us         0.16%      32.000us      32.000us       0.000us         0.00%     624.000us     624.000us           0 b           0 b           0 b           0 b             1  
                                       aten::clamp_min_         0.04%       9.000us         0.08%      16.000us      16.000us     624.000us         3.59%     624.000us     624.000us           0 b           0 b           0 b           0 b             1  
                                  cudaDeviceSynchronize        95.72%      19.748ms        95.72%      19.748ms      19.748ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             1  
void at::native::elementwise_kernel&lt;128, 2, at::nati...         0.00%       0.000us         0.00%       0.000us       0.000us       2.180ms        12.54%       2.180ms       2.180ms           0 b           0 b           0 b           0 b             1  
                                ampere_sgemm_128x128_nn         0.00%       0.000us         0.00%       0.000us       0.000us      14.583ms        83.87%      14.583ms       7.292ms           0 b           0 b           0 b           0 b             2  
void at::native::vectorized_elementwise_kernel&lt;4, at...         0.00%       0.000us         0.00%       0.000us       0.000us     624.000us         3.59%     624.000us     624.000us           0 b           0 b           0 b           0 b             1  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 20.632ms
Self CUDA time total: 17.387ms

</code></pre>

Szymon_Antoniak: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> does it look as expected? Thanks</p>