Mixed precision VQ-VAE makes NaN loss

clayton: <p>Hello,<br>
I’ve been trying to apply automatic mixed precision on this <a href="https://github.com/rosinality/vq-vae-2-pytorch" rel="noopener nofollow ugc">VQ-VAE implementation</a><br>
by following the pytorch documentation:</p>
<pre><code class="lang-auto">        with autocast():
            out, latent_loss = model(img)
            recon_loss = criterion(out, img)
            latent_loss = latent_loss.mean()
            loss = recon_loss + latent_loss_weight * latent_loss
        scaler.scale(loss).backward()

        if scheduler is not None: #not using scheduler
            scheduler.step()
        scaler.step(optimizer)
        scaler.update()
</code></pre>
<p>Unfortunately, the MSE appears once for a split second and then immediately goes to nan.<br>
I don’t know if it is possible to use AMP on a VQ-VAE so any help would be appreciated,<br>
Thanks</p>

ptrblck: <p>I assume the model converges in <code>float32</code> and the loss doesn’t blow up?<br>
If so, could you check the parameters after each <code>step</code> for invalid values (NaNs and Infs)?</p>

clayton: <p>Yes, the model converges in float32. I have tried checking for NaNs and Infs with this code:</p>
<pre><code class="lang-auto">        scaler.step(optimizer)
        for param in model.parameters():
            if torch.isnan(param.data).any() or torch.isinf(param.data).any():
            	print(param.data)
</code></pre>
<p>but it doesn’t show anything. The loss is around 0.38 after the first step and then goes to NaN beacause the tensors returned by  <code>out, latent_loss = model(img)</code>  are filled with only NaNs.<br>
The anomaly detection gives me “RuntimeError: Function ‘MseLossBackward’ returned nan values in its 0th output.”</p>
<p>I have isolated the problem and it seems to be in the Decoder since the NaNs appear after line 213<br>
<code>dec_t = self.dec_t(quant_t)</code> in <a href="https://github.com/rosinality/vq-vae-2-pytorch/blob/master/vqvae.py" rel="noopener nofollow ugc">vqvae.py</a></p>
<p>Thanks for helping</p>

ptrblck: <p>Thanks for the debugging. Could you disable <code>autocast</code> for <code>self.dec_t</code> and check the output range? This could narrow down, if the decoder outputs are indeed overflowing or if the NaNs are created in a different operation.<br>
I’m also not sure how the quantization works in this model. Could you explain it a bit how it would quantize the values for FP32?</p>

clayton: <p>Putting <code>with torch.cuda.amp.autocast(enabled=False):</code> before <code>dec_t = self.dec_t(quant_t)</code> doesn’t change anything, but if I put it in forward pass of the decoder I now get the error <code>RuntimeError: Function 'CudnnConvolutionBackward' returned nan values in its 1th output.</code> and the traceback points to the forward pass. Sorry if I’m doing this wrong, I’m not quite sure where I should disable autocast.</p>

ptrblck: <p>No, you debugging sounds alright.<br>
If you disable autocast for the decoder, were you able to isolate which method introduces the NaNs? Is it still the decoder?</p>

clayton: <p>I cannot not find any NaNs in the foward pass when disabling autocast for the decoder, but the traceback of the error says it is still the decoder that produces NaNs, more specifically a convolution layer in the decoder:</p>
<p>\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl<br>
result = self.forward(*input, **kwargs)<br>
lib\site-packages\torch\nn\modules\conv.py", line 423, in forward<br>
return self._conv_forward(input, self.weight)<br>
lib\site-packages\torch\nn\modules\conv.py", line 420, in _conv_forward<br>
self.padding, self.dilation, self.groups)<br>
(function _print_stack)</p>
<p>I have modified the Sequential block into a line of code for each layer and the traceback now points to the first layer of the decoder which is a simple convolution: <code>nn.Conv2d(in_channel, channel, 3, padding=1)</code></p>

ptrblck: <p>Could you check the inputs, output, as well as the parameters of this layer and their stats in particular?<br>
Since this layer is now running in <code>float32</code>, the values are either huge in their magnitude, or another layer is causing the invalid values.</p>

clayton: <p>I have checked the inputs and found some NaNs so I found the responsible layer which was a <code>nn.ConvTranspose2d(embed_dim, embed_dim, 4, stride=2, padding=1 )</code> so I made it run in float32 and now there’s no more NaNs but the loss explodes to values like 6527757.50000 . So again, I find the responsible layer, <code>self.quantize_t = Quantize(embed_dim, n_embed)</code> but when I try to make it run in float32 I get the error:</p>
<p>\vqvae.py", line 46, in forward<br>
+ self.embed.pow(2).sum(0, keepdim=True)<br>
RuntimeError: expected scalar type Half but found Float</p>
<p>during the forward pass of the layer.<br>
At this point, im not sure this is even feasible.</p>

ptrblck: <p>The exploding loss sounds concerning, which is why I asked about it in the first post, as it would blow up mixed-precision training.<br>
I don’t quite understand the last error. Is the <code>Quantize</code> layer depending on the usage of amp or why is it suddenly not working in <code>float32</code>?</p>

clayton: <p>I don’t really understand it either. It only happens when I try to run mixed precision with the problematic layers, including the <code>Quantize</code> layer in float32. If everything is in float32 theres no problem. Even if theres is a way to fix this I’m not sure I would get any significant speedup since most layers would have to run in float32.</p>

f10w: <p><a class="mention" href="/u/clayton">@clayton</a> Were you be able to overcome this issue? Thanks.</p>