Mixed Precision training not working with Cudnn benchmark

tyb_10: <p>Dear all,<br>
I have upgraded torch to 1.6 to use native mixed precision training.<br>
I am observing some strange behavior that mixed precision training does not seem to have any effect on model memory consumption with <code>cudnn_benchmark=True</code>.<br>
I am using pytorch lightning so i set the mixed precision training as<br>
System<br>
*Pytorch 1.6 *<br>
<em>Pytorch lightning 0.8.1</em><br>
<em>Linux 18.01</em><br>
<em>GPU Nvidia Tesla T4</em></p>
<p><code>trainer(...., precision=16, amp_level='O2')</code></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/3/4303d326d42ea506e21b0205b30593ae271c1f9b.png" data-download-href="https://discuss.pytorch.org/uploads/default/4303d326d42ea506e21b0205b30593ae271c1f9b" title="Screen Shot 2020-08-09 at 1.38.49 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/3/4303d326d42ea506e21b0205b30593ae271c1f9b_2_690x198.png" alt="Screen Shot 2020-08-09 at 1.38.49 PM" data-base62-sha1="9yQbtG86vKZWTDjch4jc47Myg9t" width="690" height="198" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/3/4303d326d42ea506e21b0205b30593ae271c1f9b_2_690x198.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/3/4303d326d42ea506e21b0205b30593ae271c1f9b_2_1035x297.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/3/4303d326d42ea506e21b0205b30593ae271c1f9b_2_1380x396.png 2x" data-dominant-color="191919"><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">Screen Shot 2020-08-09 at 1.38.49 PM</span><span class="informations">1666×480 135 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>Following is output of nvidia-smi with <strong>FP16 enabled</strong></p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.05    Driver Version: 450.51.05    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   73C    P0    64W /  70W |  12230MiB / 15109MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
</code></pre>
<p>following is <strong>with FP16 disabled</strong></p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.05    Driver Version: 450.51.05    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   77C    P0    63W /  70W |  12200MiB / 15109MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
</code></pre>
<p>I was playing around with other parameters and found out that if i disable <code>cudnn_benchmark</code><br>
I see ~ 14% memory reduction 12.2 GB → 10.5 GB but by <em>disabling cudnn_benchmark my forward pass time increases significantly.</em><br>
My model has 12 million trainable parameters.<br>
I was expecting it to reduce much further as the model is quite big so the effect of FP16 should be more evident. My model contains quite some 3D convolutions as well.</p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.05    Driver Version: 450.51.05    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   56C    P0    66W /  70W |  10578MiB / 15109MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
</code></pre>
<p>What do you think is the reason of such a behavior?<br>
thanks.</p>

ptrblck: <p>The memory usage given in <code>nvidia-smi</code> will give you the reserved memory in PyTorch (allocated + cached) as well as the CUDA context (and all other processes).<br>
Setting <code>cudnn.benchmark=True</code> will run different cudnn kernels for each new input shape and select the fastest one. This behavior might use more memory in the initial iteration and the cached memory will thus increase. In case a profiled algorithm uses a workspace, which would cause an out of memory issue, the algo is skipped.<br>
To check the allocated memory, you can use <code>torch.cuda.memory_allocated()</code>.</p>

tyb_10: <p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
sorry for the late reply,<br>
so I profiled the memory allocation at the start of the training loop and reporting values at the start of 3rd iteration after which these values remain constant for each setting.<br>
If you could please explain this behavior, it would be really helpful.<br>
p.s just for validation, I have tried <code>gc.collect()</code> , deleting input tensors, outputs, clearing cache but these strategies don’t change anything.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/5/951615a9e87b81ce21e537edec1cac457850aaad.png" data-download-href="https://discuss.pytorch.org/uploads/default/951615a9e87b81ce21e537edec1cac457850aaad" title="Screen Shot 2020-08-15 at 7.23.10 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/951615a9e87b81ce21e537edec1cac457850aaad_2_690x239.png" alt="Screen Shot 2020-08-15 at 7.23.10 PM" data-base62-sha1="lgSst6HRoV2Kiry1UidhPXxRp8h" width="690" height="239" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/951615a9e87b81ce21e537edec1cac457850aaad_2_690x239.png, https://discuss.pytorch.org/uploads/default/original/3X/9/5/951615a9e87b81ce21e537edec1cac457850aaad.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/5/951615a9e87b81ce21e537edec1cac457850aaad.png 2x" data-dominant-color="F4F4F4"><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">Screen Shot 2020-08-15 at 7.23.10 PM</span><span class="informations">987×342 29.4 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>Based on your output, the <code>FP16</code> approach uses less memory in the run, while creates a higher peak memory.<br>
As explained before, the higher memory peak might be created e.g. by the different profiled kernels when <code>cudnn.benchmark=True</code> is used.</p>

tyb_10: <p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
As we know <code>cudnn.benchmark=True</code> is generally recommended  for speed ups so turning it off would slow down the training considerably.<br>
Also, even if FP16 uses less memory in the run, we see that even with <code>cudnn</code> disabled the reduction of memory is very minimal (cause peak memory consumptions would cause the OOM) so unless and until these peak memory consumptions go down, GPUs will run out of memory. What should be the mitigation strategy in this case.</p>

ptrblck: <aside class="quote no-group" data-username="tyb_10" data-post="5" data-topic="92190">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/8baadc/48.png" class="avatar"> tyb_10:</div>
<blockquote>
<p>(cause peak memory consumptions would cause the OOM)</p>
</blockquote>
</aside>
<p>No, that shouldn’t be the case, as explained before:</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="92190">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>In case a profiled algorithm uses a workspace, which would cause an out of memory issue, the algo is skipped.</p>
</blockquote>
</aside>
<p>Could you post a setup, where an OOM error is raised?</p>