MSE Loss becomes NaN after a few iterations

Apurva: <p>I am training a deep model with an LSTM and GNNs. It is a regression problem so the loss is MSE Loss. I am seeing that the loss becomes NaN after a few iterations. I turned on torch.autograd.set_detect_anomaly, and here is what I get:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/4/c487f9411259a216bcdf2d1ab7d052ef9692078d.png" data-download-href="https://discuss.pytorch.org/uploads/default/c487f9411259a216bcdf2d1ab7d052ef9692078d" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/4/c487f9411259a216bcdf2d1ab7d052ef9692078d.png" alt="image" data-base62-sha1="s2AUcem0KTB3cI6EcHLdhRQe3Dv" width="690" height="370" data-dominant-color="212121"><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">image</span><span class="informations">1049×564 32.2 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>The code in question is:<br>
…<br>
…<br>
order = torch.argsort(lengths, descending=True)<br>
order_rev = torch.argsort(order)<br>
x, lengths = x[order], lengths[order]</p>
<pre><code>    # feed-forward bidirectional RNN
    self.rnn.flatten_parameters()
    total_length = x.shape[1]
    h = self.embed(x)
    h = pack_padded_sequence(h, lengths, batch_first=True)
    h, _ = self.rnn(h)
    h = pad_packed_sequence(h, batch_first=True, total_length=total_length)[0][:, 1:-1][order_rev]
</code></pre>
<p>…<br>
…</p>
<p>I inserted forward hooks, and checked the parameters after every training loop like so:</p>
<p>def nan_hook(self, inp, output):<br>
if type(output) == torch.nn.utils.rnn.PackedSequence:<br>
outputs = [output.data]<br>
elif not isinstance(output, tuple):<br>
if type(output[0]) == torch.nn.utils.rnn.PackedSequence:<br>
outputs = [output.data]<br>
else:<br>
outputs = [output]<br>
else:<br>
if type(output[0]) == torch.nn.utils.rnn.PackedSequence:<br>
outputs = [output[0].data]<br>
else:<br>
outputs = output</p>
<pre><code>        for i, out in enumerate(outputs):
            nan_mask = torch.isnan(out)
            if nan_mask.any():
                print("In", self.__class__.__name__)
                raise RuntimeError(f"Found NAN in output {i} at indices: ", nan_mask.nonzero(), "where:", out[nan_mask.nonzero()[:, 0].unique(sorted=True)])
            inf_mask = torch.isinf(out)
            if inf_mask.any():
                print("In", self.__class__.__name__)
                raise RuntimeError(f"Found INF in output {i} at indices: ", inf_mask.nonzero(), "where:", out[inf_mask.nonzero()[:, 0].unique(sorted=True)])
</code></pre>
<p>for submodule in model.modules():<br>
submodule.register_forward_hook(nan_hook)</p>
<p>Checking for parameters after every training loop:</p>
<pre><code>    for model in self.models_dict["model"]:
        for param in model.parameters():
            if not torch.isfinite(param.data).all():
                print("param data:", torch.isfinite(param.data).all())
</code></pre>
<p>None of these checks fire indicating that there are no NaNs in the data or in any of the parameters.</p>
<p>My machine configuration is:<br>
NVIDIA-SMI 460.27.04<br>
Driver: 460.27.04<br>
CUDA 11.2<br>
GPU: V100, 16GB</p>
<p>I am using python 3.8.5, and pytorch 1.6.0 and cudatoolkit 10.2.89</p>
<p>This problem started coming after IT upgraded the machine. I am trying to find out from IT what upgrade they did on the machine.</p>
<p>I tried upgrading to pytorch1.7 and cudatoolkit11.0, but that combination seems to be unstable and has its own share of problems.</p>
<p>I have searched on the forums, and have tried the usual techniques (reducing learning rate, adding L2 regularization etc), but those just move the problem around to a later point. Gradient clipping is also turned on. The problem never really goes away.</p>
<p>Any suggestions to debug will be very welcome.</p>
<p>Thanks,<br>
Apurva</p>