How to pinpoint NaN grads?

CesareMontresor: <p>Hello there,</p>
<p>I’m using an optimizer (AdamW) to optimize the positions of some points in a N dimentional space, based on some grouping patterns:</p>
<p>At each iteration:</p>
<ul>
<li>A subset of those points are selected, and they should move closer.</li>
<li>At the same time, the unselected points, should move far away.</li>
<li>I’m using a loss that resembles the contrastive learning loss from clip.</li>
</ul>
<p>Here an example of the performed operations:</p>
<pre data-code-wrap="py"><code class="lang-py">        bs = selected.size(0)
        
        idxs = (selected &gt;= 0)
        idxs_on = (idxs)
        idxs_off = (~idxs)

        cnt = idxs_on.sum(dim=1)
        if (cnt==0).any():
            raise ValueError('cnt has 0')

        vals_on = selected*idxs_on
        vals_off = -(selected*idxs_off)

        vals_on = (vals_on / vals_on.sum(dim=1).unsqueeze(1)).unsqueeze(2)
        vals_off = (vals_off / vals_off.sum(dim=1).unsqueeze(1)).unsqueeze(2)

        batch_positions = positions.expand( bs, -1, -1 )

        center_on = (batch_positions * vals_on).sum(dim=1).unsqueeze(1).detach()
        center_off = (batch_positions * vals_off).sum(dim=1).unsqueeze(1)
        
        center_dist = (batch_positions - center_on)

        if torch.isnan(center_dist).any():
            raise ValueError('center_dist has nan')

        #minimize
        dist_on = (center_dist*idxs_on.unsqueeze(2)).pow(2).sum(dim=1).sqrt()
        #maximize
        dist_off = (center_dist*idxs_off.unsqueeze(2)).pow(2).sum(dim=1).sqrt()
        #maximize
        dist_center = (center_on - center_off).pow(2).sum(dim=1).sqrt()

        dist_on_loss = dist_on.sum(dim=1)
        dist_off_loss = torch.log(1+dist_off).sum(dim=1)
        dist_center_loss  = dist_center.sum(dim=1)
        

        loss = dist_on_loss - dist_off_loss - dist_center_loss
        
        loss.backward()
</code></pre>
<p>( i say an example, bacause this code operates on several batches of points and several different groups, but this is the essence of it )</p>
<p>After some iterations my grads goes to NaN.</p>
<p>I’ve tried to:</p>
<ul>
<li>set a very small learning rate (1e-10)</li>
<li>play with the batchsize</li>
<li>monitor the forward pass and look for Inf, NaN and Zeros<br>
But the forward pass looks ok.</li>
</ul>
<p>I added</p>
<pre data-code-wrap="py"><code class="lang-py">torch.autograd.set_detect_anomaly(True)
</code></pre>
<p>and it gives me</p>
<pre><code class="lang-auto">RuntimeError: Function 'PowBackward0' returned nan values in its 0th output.
</code></pre>
<p>Unfortunatly it always display it on the loss.backward() line</p>
<p>Is there a way to pinpoint in which operation exactly the grads goes to NaN during the backward pass? is there a way to place a breakpoint there?<br>
Any input is welcome.</p>
<p>Thanks &lt;3</p>

CesareMontresor: <p>It turned out that in my pow(2) methods, during the forward pass, I has a bunch of near to 0 vars.</p>
<p>While 0^2, gives just a 0, a near-to-0 values, leads to all sort of numerical instabilities.</p>
<p>Still I really would like to find a way to debug NaN cards more effectively.<br>
No other way?</p>

soulitzer: <blockquote>
<p>Is there a way to pinpoint in which operation exactly the grads goes to NaN during the backward pass?</p>
</blockquote>
<p>If you scroll up the anomaly detection message should show you a stack trace of which operation during forward corresponds to the problematic PowBackward0.</p>