Layer output inconsistent on CUDA

davidz0123: <p>I encountered inconsistent Conv2d layer outputs on CUDA, although on CPUs the outputs are consistent. The differences are bigger than the floating point precision, and switching to double precision did not help. Below is an example script. When I ran it for the second time, I also received an error about illegal memory access.</p>
<pre><code class="lang-auto"># example.py

import torch
from torch.nn import Conv2d

torch.manual_seed(0)

shape = (30, 3, 256, 256)
x = torch.linspace(0, 1, shape[0] * shape[1] * shape[2] * shape[3]).view(shape)
layer = Conv2d(3, 1024, kernel_size=(16, 16), stride=(16, 16))

print('cpu:')
print(layer(x[:1]).mean((1, 2, 3)).item())
print(layer(x[:1]).mean((1, 2, 3)).item())
print(layer(x[:2]).mean((1, 2, 3))[:1].item())
print(layer(x[:2]).mean((1, 2, 3))[:1].item())

x = x.to('cuda')
layer = layer.to('cuda')

print('cuda:')
print(layer(x[:1]).mean((1, 2, 3)).item())
print(layer(x[:1]).mean((1, 2, 3)).item())
print(layer(x[:2]).mean((1, 2, 3))[:1].item())
print(layer(x[:2]).mean((1, 2, 3))[:1].item())

print('cuda again:')
print(layer(x[:1]).mean((1, 2, 3)).item())
print(layer(x[:1]).mean((1, 2, 3)).item())
print(layer(x[:2]).mean((1, 2, 3))[:1].item())
print(layer(x[:2]).mean((1, 2, 3))[:1].item())

layer = layer.double()
x = x.double()

print('cuda double precision:')
print(layer(x[:1]).mean((1, 2, 3)).item())
print(layer(x[:1]).mean((1, 2, 3)).item())
print(layer(x[:2]).mean((1, 2, 3))[:1].item())
print(layer(x[:2]).mean((1, 2, 3))[:1].item())
</code></pre>
<p>The outputs I received are:</p>
<pre><code class="lang-auto">$ python example.py  # first run
cpu:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda:
0.0011005939450114965
0.0005542933358810842
0.0005991645157337189
0.0011005939450114965
cuda again:
0.003864584257826209
0.0011005939450114965
0.0011006519198417664
0.0011010946473106742
cuda double precision:
0.0011005939325371148
0.0016997585289742778
0.0004818475426274861
0.0004950613530068367

$ python example.py  # second run
cpu:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda:
0.0011005939450114965
0.0011032642796635628
0.0005991645157337189
0.0011005939450114965
cuda again:
0.003864584257826209
0.0011005939450114965
0.0004513512831181288
0.0010505157988518476
cuda double precision:
0.0011005939325371148
Traceback (most recent call last):
  File "/home/my_username/conv2d_cuda.py", line 36, in &lt;module&gt;
    print(layer(x[:1]).mean((1, 2, 3)).item())
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
</code></pre>

ptrblck: <p>I cannot reproduce any issues on a 3090 running <code>torch==2.4.0.dev20240610+cu124</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">cpu:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda again:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda double precision:
0.0011005939325371148
0.0011005939325371148
0.0011005939325371148
0.0011005939325371148
</code></pre>
<p>In case you are using an older release, I would recommend updating PyTorch.</p>

davidz0123: <p>I am using torch 2.3.1+cu121 (as shown by <code>print(torch.__version__)</code>). I tried running the example script on both RTX 2080 Ti and A100-SXM4-80GB and encountered the same problem.</p>

ptrblck: <p>I’m also unable to reproduce the issue using <code>torch==2.3.1+cu121</code> on an A100 node and get:</p>
<pre data-code-wrap="python"><code class="lang-python">cpu:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda:
0.0011006121058017015
0.0011006121058017015
0.0011005939450114965
0.0011005939450114965
cuda again:
0.0011006121058017015
0.0011006121058017015
0.0011005939450114965
0.0011005939450114965
cuda double precision:
0.0011005939325371148
0.0011005939325371148
0.0011005939325371148
0.0011005939325371148
</code></pre>
<p>Note that the small differences (which are in the expected range) are caused by TF32 as different kernels are used if you slice the input differently, which you can disable via <code>torch.backends.cudnn.allow_tf32 = False</code>.</p>

davidz0123: <p>I tried running the script on another server, and the numbers are correct. It seems that the issue is tied to the first server I am using.</p>

davidz0123: <p>I solved the problem by setting <code>CUDA_LAUNCH_BLOCKING=1</code></p>
<pre><code class="lang-auto">$ CUDA_LAUNCH_BLOCKING=1 python example.py
cpu:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda again:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda double precision:
0.0011005939325371148
0.0011005939325371148
0.0011005939325371148
0.0011005939325371148
</code></pre>
<p>The error reappears if I set <code>CUDA_LAUNCH_BLOCKING=0</code>:</p>
<pre><code class="lang-auto">$ CUDA_LAUNCH_BLOCKING=0 python conv2d_cuda.py
cpu:
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
0.0011005939450114965
cuda:
0.0011005939450114965
0.0005014293710701168
0.0005991645157337189
0.0011005939450114965
cuda again:
0.003864584257826209
0.0011005939450114965
0.0011019300436601043
0.0011007939465343952
cuda double precision:
0.0011005939325371148
Traceback (most recent call last):
  File "/home/daiweiz/hist-embs/conv2d_cuda.py", line 36, in &lt;module&gt;
    print(layer(x[:1]).mean((1, 2, 3)).item())
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
</code></pre>
<p>I looked it up and found that <code>CUDA_LAUNCH_BLOCKING=1</code> makes CUDA API calls blocking and synchronous. This might be the cause of the problem on my server.</p>

ptrblck: <p>While this env variable seems to work as a workaround, I would recommend trying to isolate what exactly is wrong on this node, as it seems that randomly wrong outputs are created and memory violations raised.</p>