Diagnosing and fixing NaNs

PeterBleackley: <p>I’ve written a module to pool the vectors from a transformer model into a single vector</p>
<pre><code class="lang-auto">class GlobalAttentionPoolingHead(torch.nn.Module):
    
    def __init__(self,config):
        """
        Creates the layer
        Parameters
        ----------
        config : transformers.RobertaConfig
                 the configuration of the model

        Returns
        -------
        None.

        """
        size = config.hidden_size
        super(GlobalAttentionPoolingHead,self).__init__()
        self.global_projection = torch.nn.parameter.Parameter(torch.empty((size,32),
                                                                         dtype=torch.float16))
        self.local_projection =  torch.nn.parameter.Parameter(torch.empty((size,32),
                                                                         dtype=torch.float16))
        sigma = (3.0/(4.0*size))**0.25
        torch.nn.init.normal_(self.global_projection,0.0,sigma)
        torch.nn.init.normal_(self.local_projection,0.0,sigma)
        self.cosine = torch.nn.CosineSimilarity(dim=2,eps=1.0e-12)
        
    
        
    def forward(self,X,attention_mask=None):
        """
        

        Parameters
        ----------
        X : torch.Tensor
            Base model vectors to apply pooling to.
        attention_mask: tensorflow.Tensor, optional
            mask for pad values
        

        Returns
        -------
        torch.Tensor
            The pooled value.

        """
        if attention_mask is None:
            attention_mask = torch.ones_like(X)
        else:
            attention_mask = attention_mask.unsqueeze(2)
        Xa = X*attention_mask
        
        gp = torch.einsum('ijk,kl-&gt;il',Xa,self.global_projection).unsqueeze(1)
        lp = torch.einsum('ijk,kl-&gt;ijl',Xa,self.local_projection)
        
        attention = self.cosine(lp,gp)
        return torch.einsum('ij,ijk-&gt;ik',attention,Xa)
</code></pre>
<p>It’s producing NaNs. So far I’ve worked out that <code>attention</code> on the second to last line contains NaNs, but I haven’t been able to trace it back further than that. I put in a test to see if either of <code>gp</code> or <code>lp</code> contained NaNs - they seemed not to, but I’m not 100% confident that the test worked.</p>
<p>How can I troubleshoot this? I’m running it on a Kaggle notebook with 2 T4 GPUs, and am using the HuggingFace <code>accelerate</code> library.</p>

Dhia-naouali: <p>I would say it’s probably caused by the extra smol epsilon value used in cosine similarity (try smth like 1e-6 or even higher fir half precision)<br>
plus, u’re using float16, underflow is easier to be triggered)</p>
<p>(I can provide more details if needed once I access my laptop)</p>

PeterBleackley: <p>Thanks for your help. On closer examination it seems that the <code>local_projection</code> is producing zero vectors - possibly an underflow. That in turn leads to the gradient of the cosine similarity becoming NaN, and thus invalid weights thereafter. Increasing the initial weights on the <code>global_projection</code> should help.</p>

PeterBleackley: <p>After increasing <code>eps</code> to 10^-4, I found that the model worked for the first batch of training data, but had NaN weights afterwards. To trace this I added the following code to the training loop after the backpropogation step</p>
<pre><code class="lang-auto">            for (name,param) in model.named_parameters():
                if torch.any(param.isnan()).item():
                    print(name)
</code></pre>
<p>Now, the vectors produced from my model go through a number of other layers to a final binary logistic regression classifier</p>
<pre><code class="lang-auto">class Classifier(torch.nn.Module):
    def __init__(self,size):
        super(Classifier,self).__init__()
        self.logit=torch.nn.Linear(size,1,
                                  dtype=torch.float16)
        self.activation=torch.nn.Sigmoid()

    def forward(self,x):
        return self.activation(self.logit(x)).squeeze()
</code></pre>
<p>This final layer was the only one that didn’t have invalid weights after the first backpropogation, which indicates that it is the one where the invalid gradients originate. But why would that be the case?</p>

PeterBleackley: <p>I was using a Nadam optimiser, which turns out to be prone to numerical instability. Changed to Rprop.</p>