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

[aced125]

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)

[cli99]

Can you clarify how the problem relates to DeepSpeed? Or can you share the full reproduce code of the problem? Thanks.

[Mindful]

I'm not sure I understand what's going on here or exactly how this relates to deepspeed, but I think I may just have encountered the same problem. I found that increasing my batch size while using deepspeed resulted in inf loss - I changed only batch size, nothing else, so I really don't understand how this could happen - but replacing NLL loss with reduction="mean" with reduction="none" and then .mean() appears to have solved the issue.

I did not see this issue in pure Pytorch training, and I only saw it with deepspeed when I pushed to batch sizes larger than what I could handle without using deepspeed.