nn.Conv2d produces NaNs when specific value of out_channels is set

Vinc: <p>I am trying to implement <code>fp16</code> training for a VQ-GAN. The issue is: <code>NaNs</code> appear in the first batch when using gradient scaler. I tracked down the issue to a specific <code>Conv2d</code> layer. I also saved the input coming to it which can be downloaded here:</p>
<aside class="onebox googledrive" data-onebox-src="https://drive.google.com/file/d/1Wv74Bp46au4mWFhVAW9CikDejIkdsiNE/view?usp=drive_link">
  <header class="source">

      <a href="https://drive.google.com/file/d/1Wv74Bp46au4mWFhVAW9CikDejIkdsiNE/view?usp=drive_link" target="_blank" rel="noopener nofollow ugc">drive.google.com</a>
  </header>

  <article class="onebox-body">
      <a href="https://drive.google.com/file/d/1Wv74Bp46au4mWFhVAW9CikDejIkdsiNE/view?usp=drive_link" target="_blank" rel="noopener nofollow ugc"><span class="googledocs-onebox-logo g-drive-logo"></span></a>



<h3><a href="https://drive.google.com/file/d/1Wv74Bp46au4mWFhVAW9CikDejIkdsiNE/view?usp=drive_link" target="_blank" rel="noopener nofollow ugc">inf.pt</a></h3>

<p>Google Drive file.</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Code to reproduce <code>NaNs</code>:</p>
<pre><code class="lang-auto">from torch.nn import GroupNorm, SiLU, Conv2d

inf = torch.load("inf.pt").cuda()


gn1 = GroupNorm(16, 128, eps=1e-05, affine=True).cuda().half()
s = SiLU()
# out_channels: 64 -&gt; 73 NaN
conv1 = Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)).cuda().half()

o_gn1 = gn1(inf)
o_s = s(o_gn1)
conv1_out = conv1(o_s)
conv1_out.min(), conv1_out.max()

(tensor(nan, device='cuda:0', dtype=torch.float16, grad_fn=&lt;MinBackward1&gt;),
 tensor(nan, device='cuda:0', dtype=torch.float16, grad_fn=&lt;MaxBackward1&gt;))
</code></pre>
<p>The weird part is that when I set <code>out_channels</code> to anything outside the [64, 73] range, <code>Nans</code> suddenly disappear. Also when I swich back to <code>fp32</code> the problem does not occur. Am I not seeing something important? I would appreciate any help, thanks.</p>
<p>I run experiments on:<br>
CUDA: cuda_11.7<br>
Torch: 2.0.1</p>

ptrblck: <p>What’s the range of the input as well as the outputs as I guess it could overflow the valid FP16 range?</p>

Vinc: <p>That is what I initially thought, but the variance in autoencoder is well-behaved.<br>
Input to <code>conv1</code> after <code>SiLU</code>:</p>
<pre><code class="lang-auto">min: -0.2785
max: 4.5039
std: 0.5737
NaNs: 0
</code></pre>
<p>Now the stats of <code>conv1.weight</code>:</p>
<pre><code class="lang-auto">min: -0.02946
max: 0.02946
std: 0.01705
</code></pre>
<p>If I increase/decrease/change the range of <code>conv1.weight</code> it still produces <code>NaNs</code>, but only if <code>out_channels</code> is in the range of [64, 73].</p>

ptrblck: <p>Are you seeing the same issue after updating PyTorch to the latest stable or nightly release?</p>

Vinc: <p>Oh, that is weird, when I switch to newer version of PyTorch on Google Colab, training continues without any issues, thanks!</p>

ptrblck: <p>Great! Thanks for checking your code with a newer release and confirming it’s working now.</p>