KL divergence negative with AMP

dominikschmidt: <p>I’m computing the KL divergence between two categorical distributions using <code>torch.distributions.kl.kl_divergence</code>. When using AMP, the result for my particular inputs is frequently negative (&gt;40% of cases), but only very rarely (&lt;1%) when not using AMP. According to the <a href="https://pytorch.org/docs/stable/amp.html#autocast-op-reference" rel="noopener nofollow ugc">autocast-op-reference</a> however, <code>kl_div</code> should autocast to <code>float32</code> anyway. I’m wondering whether this is an oversight and only applies to <code>torch.nn.functional.kl_div</code> and not to <code>torch.distributions.kl.kl_divergence</code>?</p>

ptrblck: <p>Could you post a minimal, executable code snippet showing the usage of <code>torch.distributions.kl.kl_divergence</code> so that we could check which internal <code>dtype</code>s are used?<br>
I’m currently unsure if this particular op is causing the issues or if its input is already problematic.</p>

dominikschmidt: <pre><code class="lang-auto">from torch.distributions import Categorical, kl_divergence
from torch.nn.functional import kl_div, log_softmax
import torch

l1 = torch.tensor([[1,2,3]], device='cuda', dtype=torch.float16)
l2 = torch.tensor([[1.01,2,3]], device='cuda', dtype=torch.float16)

print('NO AMP:')
print(kl_divergence(Categorical(logits=l1), 
                    Categorical(logits=l2)).item())
print(kl_div(log_softmax(l2, dim=-1), 
             log_softmax(l1, dim=-1), 
             reduction='batchmean', log_target=True).item())
print(Categorical(logits=l1).logits.dtype, 
      log_softmax(l1, dim=-1).dtype)


with torch.cuda.amp.autocast():
    print('\nAMP:')
    print(kl_divergence(Categorical(logits=l1), 
                        Categorical(logits=l2)).item())
    print(kl_div(log_softmax(l2, dim=-1), 
                 log_softmax(l1, dim=-1), 
                 reduction='batchmean', log_target=True).item())
    print(Categorical(logits=l1).logits.dtype, 
          log_softmax(l1, dim=-1).dtype)


NO AMP:
-0.0008792877197265625
-0.00015211105346679688
torch.float16 torch.float16

AMP:
-0.0008792047156020999
3.946683136746287e-06
torch.float16 torch.float32

(correct kl-div should be ~4.1e-06)
</code></pre>
<p>You’re right. It seems that neither result is correct when the inputs are already in <code>float16</code>. When using autocast, <code>log_softmax</code> casts to <code>float32</code>, that’s why it seemed that <code>nn.functional.kl_div</code> was correct. For this example, the results from <code>functional.kl_div</code> are always better than <code>distributions.kl_divergence</code>, whether inputs are both in <code>float32</code>, <code>float16</code>, or mixed.</p>
<p>I suppose in an autocast region:</p>
<ul>
<li>
<code>Categorical</code> should compute <code>logits</code> in <code>float32</code> (like <code>log_softmax</code>)?</li>
<li>and both kl-div functions should cast inputs to <code>float32</code>?</li>
</ul>

ptrblck: <p>Thanks for the update and the great debugging!<br>
Indeed it looks strange that <code>kl_divergence(Categorical())</code> seems to be giving a worse numerical stability than the <code>kl_div</code> counterpart. We’ll check how the internal implementations are used.</p>