torch Conv1d GPU memory spikes during backward prop. on smaller inputs

username: roif
cooked: <p><em>(Short reproducible code below.)</em></p>
<p>A <strong>very</strong> bizarre behavior while using torch.Conv1d:<br>
when feeding a smaller input (below some threshold), the GPU memory usage spikes sharply during backward - an order of magnitude or more.</p>
<p>We hypothesize that this has to do with torch/cuda using different algorithms for the convolution depending on dimensions and available memory; problem is, that this results very undesired OOM errors on runtime.</p>
<pre><code>import torch
from torch import nn
%load_ext pytorch_memlab

base_ch = 512
d_pos = 64


def print_gpu_mem_usage(prefix=""):
    print(f"{prefix}Peak memory: {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f} GB"
          f" | {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB"
          f" (Current: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB"
          f" | {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB)")


def isolated_conv(v):
    samp_conv = nn.Conv1d(base_ch + d_pos, 2 * base_ch, kernel_size=1, padding='valid').cuda()
    mn = samp_conv(v).mean()
    mn.backward()
</code></pre>
<p><code>%mlrun -f isolated_conv isolated_conv(torch.rand(5000, base_ch+d_pos, 11).cuda())</code></p>
<div class="md-table">
<table>
<thead>
<tr>
<th>:active_bytes:</th>
<th>:reserved_bytes:</th>
<th>:line:</th>
<th>:code:</th>
</tr>
</thead>
<tbody>
<tr>
<td>all</td>
<td>all</td>
<td></td>
<td></td>
</tr>
<tr>
<td>----peak----</td>
<td>-----peak-----</td>
<td>----</td>
<td>----</td>
</tr>
<tr>
<td>108.00M</td>
<td>108.00</td>
<td>6</td>
<td>def isolated_conv(v):</td>
</tr>
<tr>
<td>328.00M</td>
<td>346.00M</td>
<td>7</td>
<td>mn = nn.Conv1d(…)</td>
</tr>
<tr>
<td>542.00M</td>
<td>562.00M</td>
<td>8</td>
<td>mn.backward()</td>
</tr>
</tbody>
</table>
</div><p>however, switch n samples from 5000 to 4000 and it explodes:</p>
<p><code>%mlrun -f isolated_conv isolated_conv(torch.rand(4000, base_ch+d_pos, 11).cuda())</code></p>
<div class="md-table">
<table>
<thead>
<tr>
<th>:active_bytes:</th>
<th>:reserved_bytes:</th>
<th>:line:</th>
<th>:code:</th>
</tr>
</thead>
<tbody>
<tr>
<td>all</td>
<td>all</td>
<td></td>
<td></td>
</tr>
<tr>
<td>----peak----</td>
<td>-----peak-----</td>
<td>----</td>
<td>----</td>
</tr>
<tr>
<td>86.00M</td>
<td>86.00</td>
<td>6</td>
<td>def isolated_conv(v):</td>
</tr>
<tr>
<td>260.00M</td>
<td>280.00M</td>
<td>7</td>
<td>mn = nn.Conv1d(…)</td>
</tr>
<tr>
<td>8.07G</td>
<td>8.25G</td>
<td>8</td>
<td>mn.backward()</td>
</tr>
</tbody>
</table>
</div><p>Same happens also if I test the two on opposite order.</p>
<p>edits/add-it:<br>
I tried torch empty_cache(), it doesn’t help. same for cudnn.benchmark=True/False.</p>
<p>Just read about <code>torch.backends.cudnn.deterministic=True</code> which slightly changes this behavior, still the problem occurs when I go down to 3000 samples. Going to try training with it and see if it at least doesn’t crash on OOM.</p>
<p>This runs on a docker, so if you can’t reproduce with the following versions I can share a Dockerfile.</p>
<p>torch == 2.0.1<br>
pytorch-memlab == 0.3.0<br>
Nvidia 2080Ti<br>
Driver Version: 525.105.17<br>
CUDA Version: 12.0</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="roif" data-post="1" data-topic="188745">
<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/r/b2d939/48.png" class="avatar"> roif:</div>
<blockquote>
<p>problem is, that this results very undesired OOM errors on runtime.</p>
</blockquote>
</aside>
<p>Could you post the code snippet resulting in the OOM as this should never happen since cuDNN won’t use an algorithms which requires a workspace larger than the available memory?</p>

username: roif
cooked: <p>I can share the section in which the OOM happens. I tracked it down on TB profiler and can tell that it happens on the backward of this conv1d -<br>
<code>nn.Conv1d(base_ch+D_POS, 2*base_ch, kernel_size=1, padding='valid'),</code><br>
which lead me to investigate why it happens on random batches during the run.</p>
<p>This function is part of a very large, complex graph implemented with PL/Torch/PG (which I can’t share unfortunately).</p>
<p>The input to the first, raw_conv is a long concatenated vector, 1 x in_channels x &lt;order of 10^3-10^4&gt; .<br>
The input to the second, samp_conv when it crashes is an array of shape, &lt;order of 10^3&gt; x in_channels x 11 .</p>
<blockquote>
<p>base_ch = 512<br>
d_pos = 64<br>
in_channels = 106<br>
d_segment_feature = 512</p>
<p>self.raw_conv = nn.Sequential(<br>
nn.Conv1d(in_channels=in_channels, out_channels=base_ch, kernel_size=1, padding=‘same’),<br>
nn.ReLU(inplace=True),<br>
nn.Conv1d(base_ch, base_ch, kernel_size=3, groups=base_ch // 32, padding=‘same’),<br>
nn.ReLU(inplace=True),<br>
nn.Conv1d(base_ch, base_ch, kernel_size=3, groups=base_ch // 32, padding=‘same’)<br>
)</p>
<p>self.init_segment_position = nn.Sequential(<br>
nn.Conv1d(5, d_pos // 2, kernel_size=3, padding=‘same’),<br>
nn.ReLU(inplace=True),<br>
nn.Conv1d(d_pos // 2, d_pos, kernel_size=3, padding=‘same’)<br>
)</p>
<p>self.samp_conv = nn.Sequential(<br>
nn.Conv1d(base_ch + d_pos, 2 * base_ch, kernel_size=1, padding=‘valid’),<br>
nn.Conv1d(2 * base_ch, 2 * base_ch, kernel_size=3, padding=‘valid’, groups=2 * base_ch // 32, stride=2),<br>
nn.ReLU(inplace=True),<br>
nn.Conv1d(2 * base_ch, 2 * base_ch, kernel_size=3, padding=‘valid’, groups=2 * base_ch // 32, stride=2),<br>
nn.ReLU(inplace=True),<br>
nn.Conv1d(2 * base_ch, 2 * base_ch, kernel_size=2, padding=‘valid’, groups=2 * base_ch // 64, stride=2),<br>
nn.ReLU(inplace=True),<br>
nn.Conv1d(2 * base_ch, d_segment_feature,<br>
kernel_size=1, padding=‘valid’, groups=d_segment_feature // 64),<br>
)</p>
</blockquote>

username: roif
cooked: <p>short update: adding<br>
<code>torch.backends.cudnn.deterministic=True</code><br>
doesn’t solve the problem, unfortunately …</p>

username: ptrblck
cooked: <p>Without a code snippet to reproduce the issue I won’t be able to help, unfortunately, since an OOM should never be raised by cuDNN. We are only using conv algorithms requiring a workspace that would fit into the GPU memory without causing OOM. If an algo would require more memory we are ignoring it and are never calling into the kernel and will not try to allocate the workspace.</p>

username: roif
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> voila, minimalistic example. Technical details follow.</p>
<p>This creates an OOM for me within a few iterations of the loop.</p>
<blockquote>
<p>import torch<br>
from torch import nn<br>
from numpy.random import randint</p>
<p>def isolated_conv(v):<br>
print(“FWD”)<br>
mn = nn.Conv1d(512, 1024, kernel_size=1, padding=‘valid’).cuda()(v).mean()<br>
print(“BWD”)<br>
mn.backward()<br>
print(“Done with conv”)</p>
<p>rand_int = lambda: randint(3000, 5000)<br>
for ii in range(100):<br>
n_samples = rand_int()<br>
n_array_sq = rand_int()<br>
print(f"iteration {ii} | n_samples: {n_samples} | n_array_sq: {n_array_sq}“)<br>
print(f”<span class="hashtag-raw">#1</span>, running 1st conv+bwd")<br>
isolated_conv(torch.rand(n_samples, 512, 11).cuda())<br>
print(f"<span class="hashtag-raw">#1</span>, create temp tensor")<br>
temp = torch.randn(n_array_sq, 1024).cuda()<br>
print(f"<span class="hashtag-raw">#1</span>, reshape temp")<br>
temp = temp.reshape(n_array_sq, 512, 2)<br>
print(f"<span class="hashtag-raw">#1</span>, delete temp")<br>
del temp</p>
</blockquote>
<p>Specs:<br>
This runs on a 2080Ti, in a docker derived from an official Torch docker:<br>
<code>FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime</code><br>
which has -<br>
Python 3.10.11<br>
torch                       2.0.1</p>

username: roif
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> - per your request, I sent a minimal reproducible example; any resolution?</p>