F.cross_entropy returns infs sometimes due to it summing the losses.

Hello,

F.cross_entropy returns infs sometimes due to it summing the losses.

Findings:

The following SOMETIMES returns inf loss (the default options in F.cross_entropy)

import torch.nn.functional as F
import torch as th

logits = th.randn(32, 1000)
labels = th.randint(low=0, high=1000, size=(32, ))
labels[18] = -100

loss = F.cross_entropy(logits, labels, ignore_index=-100, reduction="mean")
print(loss)

The following will NOT return inf

loss = F.cross_entropy(logits, labels, ignore_index=-100, reduction="none")
loss = loss.mean()
print(loss)

The following SOMETIMES returns inf

loss = F.cross_entropy(logits, labels, ignore_index=-100, reduction="none")
loss = loss.sum()
print(loss)

The solution I am currently using to do normal cross entropy in DeepSpeed:

loss = F.cross_entropy(logits, labels, reduction="none")
numel = labels.numel()
numel_no_mask = labels.ne(-100).sum()
norm = numel_no_mask / numel

loss = loss.mean() / norm
print(loss)

Reproduce the reported F.cross_entropy examples with ignore_index=-100 and compare the mean, none, and sum reductions. Then trace the DeepSpeed loss path referenced in the report to determine whether the correction belongs in DeepSpeed or upstream PyTorch; done should mean masked cross-entropy no longer produces inf values and the behavior is covered by a regression test.